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

alita_sdk/tools/confluence/api_wrapper.py

@@ -1,24 +1,27 @@
-import re
-import logging
-import requests
-import json
 import base64
+import json
+import logging
+import re
 import traceback
-from typing import Optional, List, Any, Dict, Callable, Generator, Literal
 from json import JSONDecodeError
+from typing import Optional, List, Any, Dict, Callable, Generator, Literal
 
-from pydantic import Field, PrivateAttr, model_validator, create_model, SecretStr
-from tenacity import retry, stop_after_attempt, wait_exponential, before_sleep_log
-
+import requests
+from atlassian.errors import ApiError
+from langchain_community.document_loaders.confluence import ContentFormat
 from langchain_core.documents import Document
-from langchain_core.tools import ToolException
 from langchain_core.messages import HumanMessage
+from langchain_core.tools import ToolException
 from markdownify import markdownify
-from langchain_community.document_loaders.confluence import ContentFormat
+from pydantic import Field, PrivateAttr, model_validator, create_model, SecretStr
+from requests import HTTPError
+from tenacity import retry, stop_after_attempt, wait_exponential, before_sleep_log
 
-from ..elitea_base import BaseVectorStoreToolApiWrapper, extend_with_vector_tools
+from alita_sdk.tools.non_code_indexer_toolkit import NonCodeIndexerToolkit
+from alita_sdk.tools.utils.available_tools_decorator import extend_with_parent_available_tools
 from ..llm.img_utils import ImageDescriptionCache
 from ..utils import is_cookie_token, parse_cookie_string
+from ...runtime.utils.utils import IndexerKeywords
 
 logger = logging.getLogger(__name__)
 
@@ -171,7 +174,7 @@ def parse_payload_params(params: Optional[str]) -> Dict[str, Any]:
     return {}
 
 
-class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
+class ConfluenceAPIWrapper(NonCodeIndexerToolkit):
     # Changed from PrivateAttr to Optional field with exclude=True
     client: Optional[Any] = Field(default=None, exclude=True)
     base_url: str
@@ -193,15 +196,7 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
     keep_markdown_format: Optional[bool] = True
     ocr_languages: Optional[str] = None
     keep_newlines: Optional[bool] = True
-    llm: Any = None
-    # indexer related
-    connection_string: Optional[SecretStr] = None
-    collection_name: Optional[str] = None
-    doctype: Optional[str] = 'doc'
-    embedding_model: Optional[str] = "HuggingFaceEmbeddings"
-    embedding_model_params: Optional[Dict[str, Any]] = {"model_name": "sentence-transformers/all-MiniLM-L6-v2"}
-    vectorstore_type: Optional[str] = "PGVector"
-
+    _errors: Optional[list[str]] = None
     _image_cache: ImageDescriptionCache = PrivateAttr(default_factory=ImageDescriptionCache)
 
     @model_validator(mode='before')
@@ -232,13 +227,13 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
         else:
             client_instance = Confluence(url=url, username=username, password=api_key, cloud=cloud)
 
-        custom_headers = values.get('custom_headers', {})
-        logger.info(f"Jira tool: custom headers length: {len(custom_headers)}")
+        custom_headers = values.get('custom_headers') or {}
+        logger.info(f"Confluence tool: custom headers length: {len(custom_headers)}")
         for header, value in custom_headers.items():
             client_instance._update_header(header, value)
 
         values['client'] = client_instance
-        return values
+        return super().validate_toolkit(values)
 
     def __unquote_confluence_space(self) -> str | None:
         if self.space:
@@ -485,28 +480,78 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
         """Gets pages with specific label in the Confluence space."""
 
         start = 0
-        pages_info = []
-        for _ in range((self.max_pages + self.limit - 1) // self.limit):
-            pages = self.client.get_all_pages_by_label(label, start=start,
-                                                       limit=self.limit)  # , expand="body.view.value"
+        pages_info: List[Dict[str, Any]] = []
+        seen_ids: set[str] = set()
+
+        # Use a while-loop driven by unique pages collected and
+        # presence of additional results instead of a fixed number
+        # of iterations based purely on max_pages/limit.
+        while len(pages_info) < (self.max_pages or 0):
+            pages = self.client.get_all_pages_by_label(
+                label,
+                start=start,
+                limit=self.limit,
+            )  # , expand="body.view.value"
             if not pages:
                 break
 
-            pages_info += [{
-                'page_id': page.metadata['id'],
-                'page_title': page.metadata['title'],
-                'page_url': page.metadata['source'],
-                'content': page.page_content
-            } for page in self.get_pages_by_id([page["id"] for page in pages])]
+            # Collect only ids we haven't processed yet to avoid
+            # calling get_page_by_id multiple times for the same
+            # Confluence page.
+            new_ids: List[str] = []
+            for p in pages:
+                page_id = p["id"] if isinstance(p, dict) else getattr(p, "id", None)
+                if page_id is None:
+                    continue
+                if page_id in seen_ids:
+                    continue
+                seen_ids.add(page_id)
+                new_ids.append(page_id)
+
+            if new_ids:
+                for page in self.get_pages_by_id(new_ids):
+                    meta = getattr(page, "metadata", {}) or {}
+                    page_id = meta.get("id")
+                    page_title = meta.get("title")
+                    page_url = meta.get("source")
+                    content = getattr(page, "page_content", None)
+
+                    if page_id is None:
+                        continue
+
+                    pages_info.append(
+                        {
+                            "page_id": page_id,
+                            "page_title": page_title,
+                            "page_url": page_url,
+                            "content": content,
+                        }
+                    )
+
+                    # Respect max_pages on unique pages collected.
+                    if len(pages_info) >= (self.max_pages or 0):
+                        break
+
+            # Advance the offset by the requested page size.
             start += self.limit
-        return pages_info
+
+            # Defensive break: if the API returns fewer items than
+            # requested, there are likely no more pages to fetch.
+            if len(pages) < self.limit:
+                break
+
+        # Slice as an extra safety net in case of any race conditions
+        # around the max_pages guard in the loop above.
+        return pages_info[: (self.max_pages or len(pages_info))]
 
     def is_public_page(self, page: dict) -> bool:
         """Check if a page is publicly accessible."""
         restrictions = self.client.get_all_restrictions_for_content(page["id"])
 
         return (
-                page["status"] == "current"
+                (page["status"] == "current"
+                # allow user to see archived content if needed
+                 or page["status"] == "archived")
                 and not restrictions["read"]["restrictions"]["user"]["results"]
                 and not restrictions["read"]["restrictions"]["group"]["results"]
         )
@@ -526,18 +571,35 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
                 ),
                 before_sleep=before_sleep_log(logger, logging.WARNING),
             )(self.client.get_page_by_id)
-            page = get_page(
-                page_id=page_id, expand=f"{self.content_format.value},version"
-            )
-            if not self.include_restricted_content and not self.is_public_page(page):
-                continue
+            try:
+                page = get_page(
+                    page_id=page_id, expand=f"{self.content_format.value},version"
+                )
+            except (ApiError, HTTPError) as e:
+                logger.error(f"Error fetching page with ID {page_id}: {e}")
+                page_content_temp = f"Confluence API Error: cannot fetch the page with ID {page_id}: {e}"
+                # store errors
+                if self._errors is None:
+                    self._errors = []
+                self._errors.append(page_content_temp)
+                return Document(page_content=page_content_temp,
+                                metadata={})
+            # TODO: update on toolkit advanced settings level as a separate feature
+            # if not self.include_restricted_content and not self.is_public_page(page):
+            #     continue
             yield self.process_page(page, skip_images)
 
+    def _log_errors(self):
+        """ Log errors encountered during toolkit execution. """
+        if self._errors:
+            logger.info(f"Errors encountered during toolkit execution: {self._errors}")
+
     def read_page_by_id(self, page_id: str, skip_images: bool = False):
         """Reads a page by its id in the Confluence space. If id is not available, but there is a title - use get_page_id first."""
         result = list(self.get_pages_by_id([page_id], skip_images))
         if not result:
-            "Page not found"
+            return f"Pages not found. Errors: {self._errors}" if self._errors \
+                else "Pages not found or you do not have access to them."
         return result[0].page_content
         # return self._strip_base64_images(result[0].page_content) if skip_images else result[0].page_content
 
@@ -547,7 +609,9 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
         :param title: title
         :param type: type of content: page or blogpost. Defaults to page
         """
-        return self.client.get_page_id(space=self.space, title=title, type=type)
+
+        result = self.client.get_page_id(space=self.space, title=title, type=type)
+        return result if result else "Page not found. Check the title or space."
 
     def _strip_base64_images(self, content):
         base64_md_pattern = r'data:image/(png|jpeg|gif);base64,[a-zA-Z0-9+/=]+'
@@ -570,7 +634,7 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
                 }
                 pages_info.append(page_info)
             start += self.limit
-        return str(pages_info)
+        return str(pages_info) if pages_info else f"Unable to find anything using query {cql}. Check space or query."
 
     def search_pages(self, query: str, skip_images: bool = False):
         """Search pages in Confluence by query text in title or page content."""
@@ -821,6 +885,10 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
         from .loader import AlitaConfluenceLoader
         from copy import copy
         content_format = kwargs.get('content_format', 'view').lower()
+
+        self._index_include_attachments = kwargs.get('include_attachments', False)
+        self._include_extensions = kwargs.get('include_extensions', [])
+        self._skip_extensions = kwargs.get('skip_extensions', [])
         base_params = {
             'url': self.base_url,
             'space_key': self.space,
@@ -852,12 +920,82 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
             yield document
 
     def _process_document(self, document: Document) -> Generator[Document, None, None]:
-        attachments = self.get_page_attachments(document.metadata.get('id'))
-        if isinstance(attachments, str):
-            logger.info(f" {document.metadata.get('id')}: {attachments}")
-            return
-        for attachment in attachments:
-            yield Document(page_content=attachment.get('content', '') or attachment.get('llm_analysis', ''), metadata=attachment.get('metadata', {}))
+        try:
+            if self._index_include_attachments:
+                page_id = document.metadata.get('id')
+                attachments = self.client.get_attachments_from_content(page_id)
+                if not attachments or not attachments.get('results'):
+                    return f"No attachments found for page ID {page_id}."
+
+                # Get attachment history for created/updated info
+                history_map = {}
+                for attachment in attachments['results']:
+                    try:
+                        hist = self.client.history(attachment['id'])
+                        history_map[attachment['id']] = hist
+                    except Exception as e:
+                        logger.warning(f"Failed to fetch history for attachment {attachment.get('title', '')}: {str(e)}")
+                        history_map[attachment['id']] = None
+
+                import re
+                for attachment in attachments['results']:
+                    title = attachment.get('title', '')
+                    file_ext = title.lower().split('.')[-1] if '.' in title else ''
+
+                    # Re-verify extension filters
+                    # Check if file should be skipped based on skip_extensions
+                    if any(re.match(re.escape(pattern).replace(r'\*', '.*') + '$', title, re.IGNORECASE)
+                           for pattern in self._skip_extensions):
+                        continue
+
+                    # Check if file should be included based on include_extensions
+                    # If include_extensions is empty, process all files (that weren't skipped)
+                    if self._include_extensions and not (
+                    any(re.match(re.escape(pattern).replace(r'\*', '.*') + '$', title, re.IGNORECASE)
+                        for pattern in self._include_extensions)):
+                        continue
+
+                    media_type = attachment.get('metadata', {}).get('mediaType', '')
+                    # Core metadata extraction with history
+                    hist = history_map.get(attachment['id']) or {}
+                    created_by = hist.get('createdBy', {}).get('displayName', '') if hist else attachment.get('creator', {}).get('displayName', '')
+                    created_date = hist.get('createdDate', '') if hist else attachment.get('created', '')
+                    last_updated = hist.get('lastUpdated', {}).get('when', '') if hist else ''
+
+                    metadata = {
+                        'name': title,
+                        'size': attachment.get('extensions', {}).get('fileSize', None),
+                        'creator': created_by,
+                        'created': created_date,
+                        'updated': last_updated,
+                        'media_type': media_type,
+                        'labels': [label['name'] for label in
+                                   attachment.get('metadata', {}).get('labels', {}).get('results', [])],
+                        'download_url': self.base_url.rstrip('/') + attachment['_links']['download'] if attachment.get(
+                            '_links', {}).get('download') else None
+                    }
+
+                    download_url = self.base_url.rstrip('/') + attachment['_links']['download']
+
+                    try:
+                        resp = self.client.request(method="GET", path=download_url[len(self.base_url):], advanced_mode=True)
+                        if resp.status_code == 200:
+                            content = resp.content
+                        else:
+                            content = f"[Failed to download {download_url}: HTTP status code {resp.status_code}]"
+                    except Exception as e:
+                        content = f"[Error downloading content: {str(e)}]"
+
+                    if isinstance(content, str):
+                        yield Document(page_content=content, metadata=metadata)
+                    else:
+                        yield Document(page_content="", metadata={
+                            **metadata,
+                            IndexerKeywords.CONTENT_FILE_NAME.value: f".{file_ext}",
+                            IndexerKeywords.CONTENT_IN_BYTES.value: content
+                        })
+        except Exception as e:
+            yield from ()
 
     def _download_image(self, image_url):
         """
@@ -1598,18 +1736,24 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
             "include_restricted_content": (Optional[bool], Field(description="Include restricted content.", default=False)),
             "include_archived_content": (Optional[bool], Field(description="Include archived content.", default=False)),
             "include_attachments": (Optional[bool], Field(description="Include attachments.", default=False)),
+            'include_extensions': (Optional[List[str]], Field(
+                description="List of file extensions to include when processing attachments: i.e. ['*.png', '*.jpg']. "
+                            "If empty, all files will be processed (except skip_extensions).",
+                default=[])),
+            'skip_extensions': (Optional[List[str]], Field(
+                description="List of file extensions to skip when processing attachments: i.e. ['*.png', '*.jpg']",
+                default=[])),
             "include_comments": (Optional[bool], Field(description="Include comments.", default=False)),
-            "include_labels": (Optional[bool], Field(description="Include labels.", default=True)),
+            "include_labels": (Optional[bool], Field(description="Include labels.", default=False)),
             "ocr_languages": (Optional[str], Field(description="OCR languages for processing attachments.", default='eng')),
             "keep_markdown_format": (Optional[bool], Field(description="Keep the markdown format.", default=True)),
             "keep_newlines": (Optional[bool], Field(description="Keep newlines in the content.", default=True)),
             "bins_with_llm": (Optional[bool], Field(description="Use LLM for processing binary files.", default=False)),
         }
 
-    @extend_with_vector_tools
+    @extend_with_parent_available_tools
     def get_available_tools(self):
-        # Confluence-specific tools
-        confluence_tools = [
+        return [
             {
                 "name": "create_page",
                 "ref": self.create_page,
@@ -1726,7 +1870,3 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
             }
         ]
 
-        # Add standardized vector search tools from base class
-        vector_search_tools = self._get_vector_search_tools()
-
-        return confluence_tools + vector_search_tools

alita_sdk/tools/confluence/loader.py

@@ -3,6 +3,7 @@ from typing import Optional, List
 from logging import getLogger
 
 import requests
+from langchain_core.documents import Document
 
 logger = getLogger(__name__)
 from PIL import Image
@@ -193,6 +194,15 @@ class AlitaConfluenceLoader(ConfluenceLoader):
         else:
             return super().process_image(link, ocr_languages)
 
+    def process_page(self, page: dict, include_attachments: bool, include_comments: bool, include_labels: bool,
+                     content_format: ContentFormat, ocr_languages: Optional[str] = None,
+                     keep_markdown_format: Optional[bool] = False, keep_newlines: bool = False) -> Document:
+        if not page.get("title"):
+            # if 'include_restricted_content' set to True, draft pages are loaded and can have no title
+            page["title"] = "Untitled"
+        return super().process_page(page, include_attachments, include_comments, include_labels, content_format,
+                                    ocr_languages, keep_markdown_format, keep_newlines)
+
     # TODO review usage
     # def process_svg(
     #         self,

alita_sdk/tools/custom_open_api/__init__.py

@@ -5,7 +5,7 @@ from pydantic import create_model, BaseModel, ConfigDict, Field
 
 from .api_wrapper import OpenApiWrapper
 from ..base.tool import BaseAction
-from ..utils import clean_string, TOOLKIT_SPLITTER
+from ..utils import clean_string
 
 name = "openapi"
 
@@ -43,14 +43,19 @@ class OpenApiToolkit(BaseToolkit):
         openapi_api_wrapper = OpenApiWrapper(**kwargs)
         available_tools = openapi_api_wrapper.get_available_tools()
         tools = []
-        prefix = clean_string(toolkit_name + TOOLKIT_SPLITTER) if toolkit_name else ''
+        # Use clean toolkit name for context (max 1000 chars in description)
+        toolkit_context = f" [Toolkit: {clean_string(toolkit_name)}]" if toolkit_name else ''
         for tool in available_tools:
             if selected_tools and tool["name"] not in selected_tools:
                 continue
+            # Add toolkit context to description with character limit
+            description = tool["description"]
+            if toolkit_context and len(description + toolkit_context) <= 1000:
+                description = description + toolkit_context
             tools.append(BaseAction(
                 api_wrapper=openapi_api_wrapper,
-                name=prefix + tool["name"],
-                description=tool["description"],
+                name=tool["name"],
+                description=description,
                 args_schema=tool["args_schema"]
             ))
         return cls(tools=tools)

alita_sdk/tools/elastic/__init__.py

@@ -5,7 +5,7 @@ from pydantic import BaseModel, ConfigDict, create_model, Field, SecretStr
 
 from .api_wrapper import ELITEAElasticApiWrapper
 from ..base.tool import BaseAction
-from ..utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
+from ..utils import clean_string, get_max_toolkit_length
 
 name = "elastic"
 
@@ -19,15 +19,13 @@ def get_tools(tool):
 
 class ElasticToolkit(BaseToolkit):
     tools: list[BaseTool] = []
-    toolkit_max_length: int = 0
 
     @staticmethod
     def toolkit_config_schema() -> BaseModel:
         selected_tools = {x['name']: x['args_schema'].schema() for x in ELITEAElasticApiWrapper.model_construct().get_available_tools()}
-        ElasticToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
         return create_model(
             name,
-            url=(str, Field(default=None, title="Elasticsearch URL", description="Elasticsearch URL", json_schema_extra={'toolkit_name': True, 'max_toolkit_length': ElasticToolkit.toolkit_max_length})),
+            url=(str, Field(default=None, title="Elasticsearch URL", description="Elasticsearch URL", json_schema_extra={'toolkit_name': True})),
             api_key=(
                 Optional[SecretStr],
                 Field(
@@ -48,14 +46,17 @@ class ElasticToolkit(BaseToolkit):
         elastic_api_wrapper = ELITEAElasticApiWrapper(**kwargs)
         available_tools = elastic_api_wrapper.get_available_tools()
         tools = []
-        prefix = clean_string(toolkit_name, ElasticToolkit.toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
         for tool in available_tools:
             if selected_tools and tool["name"] not in selected_tools:
                 continue
+            description = tool["description"]
+            if toolkit_name:
+                description = f"Toolkit: {toolkit_name}\n{description}"
+            description = description[:1000]
             tools.append(BaseAction(
                 api_wrapper=elastic_api_wrapper,
-                name=prefix + tool["name"],
-                description=tool["description"],
+                name=tool["name"],
+                description=description,
                 args_schema=tool["args_schema"]
             ))
         return cls(tools=tools)