alita-sdk 0.3.263__py3-none-any.whl → 0.3.499__py3-none-any.whl

alita_sdk/runtime/clients/client.py

@@ -20,8 +20,9 @@ from .prompt import AlitaPrompt
 from .datasource import AlitaDataSource
 from .artifact import Artifact
 from ..langchain.chat_message_template import Jinja2TemplatedChatMessagesTemplate
-from ..utils.utils import TOOLKIT_SPLITTER
+from ..utils.mcp_oauth import McpAuthorizationRequired
 from ...tools import get_available_toolkit_models
+from ...tools.base_indexer_toolkit import IndexTools
 
 logger = logging.getLogger(__name__)
 
@@ -68,7 +69,11 @@ class AlitaClient:
         self.bucket_url = f"{self.base_url}{self.api_path}/artifacts/buckets/{self.project_id}"
         self.configurations_url = f'{self.base_url}{self.api_path}/integrations/integrations/default/{self.project_id}?section=configurations&unsecret=true'
         self.ai_section_url = f'{self.base_url}{self.api_path}/integrations/integrations/default/{self.project_id}?section=ai'
+        self.models_url = f'{self.base_url}{self.api_path}/configurations/models/{self.project_id}?include_shared=true'
+        self.image_generation_url = f"{self.base_url}{self.llm_path}/images/generations"
         self.configurations: list = configurations or []
+        self.model_timeout = kwargs.get('model_timeout', 120)
+        self.model_image_generation = kwargs.get('model_image_generation')
 
     def get_mcp_toolkits(self):
         if user_id := self._get_real_user_id():
@@ -81,8 +86,23 @@ class AlitaClient:
     def mcp_tool_call(self, params: dict[str, Any]):
         if user_id := self._get_real_user_id():
             url = f"{self.mcp_tools_call}/{user_id}"
-            data = requests.post(url, headers=self.headers, json=params, verify=False).json()
-            return data
+            #
+            # This loop iterates over each key-value pair in the arguments dictionary,
+            # and if a value is a Pydantic object, it replaces it with its dictionary representation using .dict().
+            for arg_name, arg_value in params.get('params', {}).get('arguments', {}).items():
+                if isinstance(arg_value, list):
+                    params['params']['arguments'][arg_name] = [
+                        item.dict() if hasattr(item, "dict") and callable(item.dict) else item
+                        for item in arg_value
+                    ]
+                elif hasattr(arg_value, "dict") and callable(arg_value.dict):
+                    params['params']['arguments'][arg_name] = arg_value.dict()
+            #
+            response = requests.post(url, headers=self.headers, json=params, verify=False)
+            try:
+                return response.json()
+            except (ValueError, TypeError):
+                return response.text
         else:
             return f"Error: Could not determine user ID for MCP tool call"
 
@@ -157,6 +177,20 @@ class AlitaClient:
             return resp.json()
         return []
 
+    def get_available_models(self):
+        """Get list of available models from the configurations API.
+
+        Returns:
+            List of model dictionaries with 'name' and other properties,
+            or empty list if request fails.
+        """
+        resp = requests.get(self.models_url, headers=self.headers, verify=False)
+        if resp.ok:
+            data = resp.json()
+            # API returns {"items": [...], ...}
+            return data.get('items', [])
+        return []
+
     def get_embeddings(self, embedding_model: str) -> OpenAIEmbeddings:
         """
         Get an instance of OpenAIEmbeddings configured with the project ID and auth token.
@@ -169,6 +203,7 @@ class AlitaClient:
             model=embedding_model,
             api_key=self.auth_token,
             openai_organization=str(self.project_id),
+            request_timeout=self.model_timeout
         )
 
     def get_llm(self, model_name: str, model_config: dict) -> ChatOpenAI:
@@ -187,20 +222,99 @@ class AlitaClient:
 
         logger.info(f"Creating ChatOpenAI model: {model_name} with config: {model_config}")
 
-        return ChatOpenAI(
-            base_url=f"{self.base_url}{self.llm_path}",
-            model=model_name,
-            api_key=self.auth_token,
-            streaming=model_config.get("streaming", True),
-            stream_usage=model_config.get("stream_usage", True),
-            max_tokens=model_config.get("max_tokens", None),
-            top_p=model_config.get("top_p"),
-            temperature=model_config.get("temperature"),
-            max_retries=model_config.get("max_retries", 3),
-            seed=model_config.get("seed", None),
-            openai_organization=str(self.project_id),
-        )
+        try:
+            from tools import this  # pylint: disable=E0401,C0415
+            worker_config = this.for_module("indexer_worker").descriptor.config
+        except:  # pylint: disable=W0702
+            worker_config = {}
+
+        use_responses_api = False
+
+        if worker_config and isinstance(worker_config, dict):
+            for target_name_tag in worker_config.get("use_responses_api_for", []):
+                if target_name_tag in model_name:
+                    use_responses_api = True
+                    break
+
+        # handle case when max_tokens are auto-configurable == -1
+        llm_max_tokens = model_config.get("max_tokens", None)
+        if llm_max_tokens and llm_max_tokens == -1:
+            logger.warning(f'User selected `MAX COMPLETION TOKENS` as `auto`')
+            # default nuber for a case when auto is selected for an agent
+            llm_max_tokens = 4000
+
+        target_kwargs = {
+            "base_url": f"{self.base_url}{self.llm_path}",
+            "model": model_name,
+            "api_key": self.auth_token,
+            "streaming": model_config.get("streaming", True),
+            "stream_usage": model_config.get("stream_usage", True),
+            "max_tokens": llm_max_tokens,
+            "temperature": model_config.get("temperature"),
+            "reasoning_effort": model_config.get("reasoning_effort"),
+            "max_retries": model_config.get("max_retries", 3),
+            "seed": model_config.get("seed", None),
+            "openai_organization": str(self.project_id),
+        }
+
+        if use_responses_api:
+            target_kwargs["use_responses_api"] = True
+
+        return ChatOpenAI(**target_kwargs)
+
+    def generate_image(self,
+                       prompt: str,
+                       n: int = 1,
+                       size: str = "auto",
+                       quality: str = "auto",
+                       response_format: str = "b64_json",
+                       style: Optional[str] = None) -> dict:
+
+        if not self.model_image_generation:
+            raise ValueError("Image generation model is not configured for this client")
+
+        image_generation_data = {
+            "prompt": prompt,
+            "model": self.model_image_generation,
+            "n": n,
+            "response_format": response_format,
+        }
+
+        # Only add optional parameters if they have meaningful values
+        if size and size.lower() != "auto":
+            image_generation_data["size"] = size
+
+        if quality and quality.lower() != "auto":
+            image_generation_data["quality"] = quality
+
+        if style:
+            image_generation_data["style"] = style
+
+        # Standard headers for image generation
+        image_headers = self.headers.copy()
+        image_headers.update({
+            "Content-Type": "application/json",
+        })
 
+        logger.info(f"Generating image with model: {self.model_image_generation}, prompt: {prompt[:50]}...")
+
+        try:
+            response = requests.post(
+                self.image_generation_url,
+                headers=image_headers,
+                json=image_generation_data,
+                verify=False,
+                timeout=self.model_timeout
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except requests.exceptions.HTTPError as e:
+            logger.error(f"Image generation failed: {e.response.status_code} - {e.response.text}")
+            raise
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Image generation request failed: {e}")
+            raise
 
     def get_app_version_details(self, application_id: int, application_version_id: int) -> dict:
         url = f"{self.application_versions}/{application_id}/{application_version_id}"
@@ -232,7 +346,8 @@ class AlitaClient:
                     app_type=None, memory=None, runtime='langchain',
                     application_variables: Optional[dict] = None,
                     version_details: Optional[dict] = None, store: Optional[BaseStore] = None,
-                    llm: Optional[ChatOpenAI] = None):
+                    llm: Optional[ChatOpenAI] = None, mcp_tokens: Optional[dict] = None,
+                    conversation_id: Optional[str] = None):
         if tools is None:
             tools = []
         if chat_history is None:
@@ -252,11 +367,15 @@ class AlitaClient:
                 if var['name'] in application_variables:
                     var.update(application_variables[var['name']])
         if llm is None:
+            max_tokens = data['llm_settings'].get('max_tokens', 4000)
+            if max_tokens == -1:
+                # default nuber for case when auto is selected for agent
+                max_tokens = 4000
             llm = self.get_llm(
                 model_name=data['llm_settings']['model_name'],
                 model_config={
-                    "max_tokens": data['llm_settings']['max_tokens'],
-                    "top_p": data['llm_settings']['top_p'],
+                    "max_tokens": max_tokens,
+                    "reasoning_effort": data['llm_settings'].get('reasoning_effort'),
                     "temperature": data['llm_settings']['temperature'],
                     "model_project_id": data['llm_settings'].get('model_project_id'),
                 }
@@ -271,13 +390,18 @@ class AlitaClient:
             app_type = "react"
         elif app_type == 'autogen':
             app_type = "react"
+
+        # LangChainAssistant constructor calls get_tools() which may raise McpAuthorizationRequired
+        # The exception will propagate naturally to the indexer worker's outer handler
         if runtime == 'nonrunnable':
             return LangChainAssistant(self, data, llm, chat_history, app_type,
-                                      tools=tools, memory=memory, store=store)
+                                      tools=tools, memory=memory, store=store, mcp_tokens=mcp_tokens,
+                                      conversation_id=conversation_id)
         if runtime == 'langchain':
             return LangChainAssistant(self, data, llm,
                                       chat_history, app_type,
-                                      tools=tools, memory=memory, store=store).runnable()
+                                      tools=tools, memory=memory, store=store, mcp_tokens=mcp_tokens,
+                                      conversation_id=conversation_id).runnable()
         elif runtime == 'llama':
             raise NotImplementedError("LLama runtime is not supported")
 
@@ -339,16 +463,50 @@ class AlitaClient:
         return self._process_requst(resp)
 
     def list_artifacts(self, bucket_name: str):
-        url = f'{self.artifacts_url}/{bucket_name}'
+        # Ensure bucket name is lowercase as required by the API
+        url = f'{self.artifacts_url}/{bucket_name.lower()}'
         data = requests.get(url, headers=self.headers, verify=False)
         return self._process_requst(data)
 
     def create_artifact(self, bucket_name, artifact_name, artifact_data):
+        # Sanitize filename to prevent regex errors during indexing
+        sanitized_name, was_modified = self._sanitize_artifact_name(artifact_name)
+        if was_modified:
+            logger.warning(f"Artifact filename sanitized: '{artifact_name}' -> '{sanitized_name}'")
+        
         url = f'{self.artifacts_url}/{bucket_name.lower()}'
         data = requests.post(url, headers=self.headers, files={
-            'file': (artifact_name, artifact_data)
+            'file': (sanitized_name, artifact_data)
         }, verify=False)
         return self._process_requst(data)
+    
+    @staticmethod
+    def _sanitize_artifact_name(filename: str) -> tuple:
+        """Sanitize filename for safe storage and regex pattern matching."""
+        import re
+        from pathlib import Path
+        
+        if not filename or not filename.strip():
+            return "unnamed_file", True
+        
+        original = filename
+        path_obj = Path(filename)
+        name = path_obj.stem
+        extension = path_obj.suffix
+        
+        # Whitelist: alphanumeric, underscore, hyphen, space, Unicode letters/digits
+        sanitized_name = re.sub(r'[^\w\s-]', '', name, flags=re.UNICODE)
+        sanitized_name = re.sub(r'[-\s]+', '-', sanitized_name)
+        sanitized_name = sanitized_name.strip('-').strip()
+        
+        if not sanitized_name:
+            sanitized_name = "file"
+        
+        if extension:
+            extension = re.sub(r'[^\w.-]', '', extension, flags=re.UNICODE)
+        
+        sanitized = sanitized_name + extension
+        return sanitized, (sanitized != original)
 
     def download_artifact(self, bucket_name, artifact_name):
         url = f'{self.artifact_url}/{bucket_name.lower()}/{artifact_name}'
@@ -490,25 +648,29 @@ class AlitaClient:
             monitoring_meta = tasknode_task.meta.get("monitoring", {})
             return monitoring_meta["user_id"]
         except Exception as e:
-            logger.warning(f"Error: Could not determine user ID for MCP tool: {e}")
+            logger.debug(f"Error: Could not determine user ID for MCP tool: {e}")
             return None
 
     def predict_agent(self, llm: ChatOpenAI, instructions: str = "You are a helpful assistant.",
                       tools: Optional[list] = None, chat_history: Optional[List[Any]] = None,
                       memory=None, runtime='langchain', variables: Optional[list] = None,
-                      store: Optional[BaseStore] = None):
+                      store: Optional[BaseStore] = None, debug_mode: Optional[bool] = False,
+                      mcp_tokens: Optional[dict] = None, conversation_id: Optional[str] = None):
         """
         Create a predict-type agent with minimal configuration.
 
         Args:
             llm: The LLM to use
             instructions: System instructions for the agent
-            tools: Optional list of tools to provide to the agent
+            tools: Optional list of tool configurations (not tool instances) to provide to the agent.
+                   Tool configs will be processed through get_tools() to create tool instances.
+                   Each tool config should have 'type', 'settings', etc.
             chat_history: Optional chat history
             memory: Optional memory/checkpointer
             runtime: Runtime type (default: 'langchain')
             variables: Optional list of variables for the agent
             store: Optional store for memory
+            debug_mode: Enable debug mode for cases when assistant can be initialized without tools
 
         Returns:
             Runnable agent ready for execution
@@ -522,17 +684,32 @@ class AlitaClient:
 
         # Create a minimal data structure for predict agent
         # All LLM settings are taken from the passed client instance
+        # Note: 'tools' here are tool CONFIGURATIONS, not tool instances
+        # They will be converted to tool instances by LangChainAssistant via get_tools()
         agent_data = {
             'instructions': instructions,
-            'tools': tools,  # Tools are handled separately in predict agents
+            'tools': tools,  # Tool configs that will be processed by get_tools()
             'variables': variables
         }
-        return LangChainAssistant(self, agent_data, llm,
-                                  chat_history, "predict", memory=memory, store=store).runnable()
+
+        # LangChainAssistant constructor calls get_tools() which may raise McpAuthorizationRequired
+        # The exception will propagate naturally to the indexer worker's outer handler
+        return LangChainAssistant(
+            self,
+            agent_data,
+            llm,
+            chat_history,
+            "predict",
+            memory=memory,
+            store=store,
+            debug_mode=debug_mode,
+            mcp_tokens=mcp_tokens,
+            conversation_id=conversation_id
+        ).runnable()
 
     def test_toolkit_tool(self, toolkit_config: dict, tool_name: str, tool_params: dict = None,
                           runtime_config: dict = None, llm_model: str = None,
-                          llm_config: dict = None) -> dict:
+                          llm_config: dict = None, mcp_tokens: dict = None) -> dict:
         """
         Test a single tool from a toolkit with given parameters and runtime callbacks.
 
@@ -551,6 +728,7 @@ class AlitaClient:
                 - configurable: Additional configuration parameters
                 - tags: Tags for the execution
             llm_model: Name of the LLM model to use (default: 'gpt-4o-mini')
+            mcp_tokens: Optional dictionary of MCP OAuth tokens by server URL
             llm_config: Configuration for the LLM containing:
                 - max_tokens: Maximum tokens for response (default: 1000)
                 - temperature: Temperature for response generation (default: 0.1)
@@ -598,7 +776,6 @@ class AlitaClient:
             llm_config = {
                 'max_tokens': 1024,
                 'temperature': 0.1,
-                'top_p': 1.0
             }
         import logging
         logger = logging.getLogger(__name__)
@@ -669,7 +846,23 @@ class AlitaClient:
                 }
 
             # Instantiate the toolkit with client and LLM support
-            tools = instantiate_toolkit_with_client(toolkit_config, llm, self)
+            try:
+                tools = instantiate_toolkit_with_client(toolkit_config, llm, self, mcp_tokens=mcp_tokens, use_prefix=False)
+            except McpAuthorizationRequired:
+                # Re-raise McpAuthorizationRequired to allow proper handling upstream
+                logger.info(f"McpAuthorizationRequired detected, re-raising")
+                raise
+            except Exception as toolkit_error:
+                # For other errors, return error response
+                return {
+                    "success": False,
+                    "error": f"Failed to instantiate toolkit '{toolkit_config.get('toolkit_name')}': {str(toolkit_error)}",
+                    "tool_name": tool_name,
+                    "toolkit_config": toolkit_config_parsed_json,
+                    "llm_model": llm_model,
+                    "events_dispatched": events_dispatched,
+                    "execution_time_seconds": 0.0
+                }
 
             if not tools:
                 return {
@@ -745,7 +938,6 @@ class AlitaClient:
             if target_tool is None:
                 available_tools = []
                 base_available_tools = []
-                full_available_tools = []
 
                 for tool in tools:
                     tool_name_attr = None
@@ -762,16 +954,14 @@ class AlitaClient:
                         if base_name not in base_available_tools:
                             base_available_tools.append(base_name)
 
-                        # Track full names separately
-                        if TOOLKIT_SPLITTER in tool_name_attr:
-                            full_available_tools.append(tool_name_attr)
-
                 # Create comprehensive error message
-                error_msg = f"Tool '{tool_name}' not found in toolkit '{toolkit_config.get('toolkit_name')}'."
+                error_msg = f"Tool '{tool_name}' not found in toolkit '{toolkit_config.get('toolkit_name')}'.\n"
+
+                # Custom error for index tools
+                if toolkit_name in [tool.value for tool in IndexTools]:
+                    error_msg += f" Please make sure proper PGVector configuration and embedding model are set in the platform.\n"
 
-                if base_available_tools and full_available_tools:
-                    error_msg += f" Available tools: {base_available_tools} (base names) or {full_available_tools} (full names)"
-                elif base_available_tools:
+                if base_available_tools:
                     error_msg += f" Available tools: {base_available_tools}"
                 elif available_tools:
                     error_msg += f" Available tools: {available_tools}"
@@ -780,10 +970,7 @@ class AlitaClient:
 
                 # Add helpful hint about naming conventions
                 if '___' in tool_name:
-                    error_msg += f" Note: You provided a full name '{tool_name}'. Try using just the base name '{extract_base_tool_name(tool_name)}'."
-                elif full_available_tools:
-                    possible_full_name = create_full_tool_name(tool_name, toolkit_name)
-                    error_msg += f" Note: You provided a base name '{tool_name}'. The full name might be '{possible_full_name}'."
+                    error_msg += f" Note: Tool names no longer use '___' prefixes. Try using just the base name '{extract_base_tool_name(tool_name)}'."
 
                 return {
                     "success": False,
@@ -889,6 +1076,9 @@ class AlitaClient:
                 }
 
         except Exception as e:
+            # Re-raise McpAuthorizationRequired to allow proper handling upstream
+            if isinstance(e, McpAuthorizationRequired):
+                raise
             logger = logging.getLogger(__name__)
             logger.error(f"Error in test_toolkit_tool: {str(e)}")
             return {
@@ -900,24 +1090,3 @@ class AlitaClient:
                 "events_dispatched": [],
                 "execution_time_seconds": 0.0
             }
-
-    def _get_real_user_id(self) -> str:
-        """Extract the real user ID from the auth token for MCP tool calls."""
-        try:
-            import base64
-            import json
-            # Assuming JWT token, extract user ID from payload
-            # This is a basic implementation - adjust based on your token format
-            token_parts = self.auth_token.split('.')
-            if len(token_parts) >= 2:
-                payload_part = token_parts[1]
-                # Add padding if needed
-                padding = len(payload_part) % 4
-                if padding:
-                    payload_part += '=' * (4 - padding)
-                payload = json.loads(base64.b64decode(payload_part))
-                return payload.get('user_id') or payload.get('sub') or payload.get('uid')
-        except Exception as e:
-            logger.error(f"Error extracting user ID from token: {e}")
-        return None
-