alita-sdk 0.3.208__py3-none-any.whl → 0.3.210__py3-none-any.whl

alita_sdk/tools/elitea_base.py

@@ -2,12 +2,16 @@ import ast
 import fnmatch
 import logging
 import traceback
-from typing import Any, Optional, List, Dict, Generator
+from typing import Any, Optional, List, Literal, Dict, Generator
 
 from langchain_core.documents import Document
 from langchain_core.tools import ToolException
-from pydantic import BaseModel, create_model, Field
+from pydantic import BaseModel, create_model, Field, SecretStr
+
+from alita_sdk.runtime.langchain.interfaces.llm_processor import get_embeddings
+from .chunkers import markdown_chunker
 from .utils import TOOLKIT_SPLITTER
+from ..runtime.utils.utils import IndexerKeywords
 
 logger = logging.getLogger(__name__)
 
@@ -98,6 +102,17 @@ BaseStepbackSearchParams = create_model(
     )),
 )
 
+BaseIndexDataParams = create_model(
+    "indexData",
+    __base__=BaseIndexParams,
+    progress_step=(Optional[int], Field(default=None, ge=0, le=100,
+                         description="Optional step size for progress reporting during indexing")),
+    clean_index=(Optional[bool], Field(default=False,
+                       description="Optional flag to enforce clean existing index before indexing new data")),
+    chunking_tool=(Literal['markdown', 'statistical', 'proposal'], Field(description="Name of chunking tool", default="markdown")),
+    chunking_config=(Optional[dict], Field(description="Chunking tool configuration", default_factory=dict)),
+)
+
 
 class BaseToolApiWrapper(BaseModel):
 
@@ -191,12 +206,34 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
 
     doctype: str = "document"
 
-    def _base_loader(self, **kwargs) -> List[Document]:
+    llm: Any = None
+    connection_string: Optional[SecretStr] = None
+    collection_name: Optional[str] = None
+    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"
+
+    def _index_tool_params(self, **kwargs) -> dict[str, tuple[type, Field]]:
+        """
+        Returns a list of fields for index_data args schema.
+        NOTE: override this method in subclasses to provide specific parameters for certain toolkit.
+        """
+        return {}
+
+    def _get_dependencies_chunker(self, document: Optional[Document] = None):
+        return markdown_chunker
+
+    def _get_dependencies_chunker_config(self, document: Optional[Document] = None):
+        embedding = get_embeddings(self.embedding_model, self.embedding_model_params)
+        #
+        return {'embedding': embedding, 'llm': self.llm}
+
+    def _base_loader(self, **kwargs) -> Generator[Document, None, None]:
         """ Loads documents from a source, processes them,
         and returns a list of Document objects with base metadata: id and created_on."""
         pass
 
-    def _process_document(self, base_document: Document) -> Document:
+    def _process_document(self, base_document: Document) -> Generator[Document, None, None]:
         """ Process an existing base document to extract relevant metadata for full document preparation.
         Used for late processing of documents after we ensure that the document has to be indexed to avoid
         time-consuming operations for documents which might be useless.
@@ -208,6 +245,51 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
             Document: The processed document with metadata."""
         pass
 
+    def get_index_data_tool(self):
+        return {
+            "name": "index_data",
+            "ref": self.index_data,
+            "description": "Loads data to index.",
+            "args_schema": create_model(
+                "IndexData",
+                __base__=BaseIndexDataParams,
+                **self._index_tool_params() if self._index_tool_params() else {}
+            )
+        }
+
+    def index_data(self, **kwargs):
+        from alita_sdk.tools.chunkers import __confluence_chunkers__ as chunkers, __confluence_models__ as models
+        docs = self._base_loader(**kwargs)
+        embedding = get_embeddings(self.embedding_model, self.embedding_model_params)
+        chunking_tool = kwargs.get("chunking_tool")
+        if chunking_tool:
+            # Resolve chunker from the provided chunking_tool name
+            base_chunker = chunkers.get(chunking_tool)
+            # Resolve chunking configuration
+            base_chunking_config = kwargs.get("chunking_config", {})
+            config_model = models.get(chunking_tool)
+            # Set required fields that should come from the instance (and Fallback for chunkers without models)
+            base_chunking_config['embedding'] = embedding
+            base_chunking_config['llm'] = self.llm
+            #
+            if config_model:
+                try:
+                    # Validate the configuration using the appropriate Pydantic model
+                    validated_config = config_model(**base_chunking_config)
+                    base_chunking_config = validated_config.model_dump()
+                except Exception as e:
+                    logger.error(f"Invalid chunking configuration for {chunking_tool}: {e}")
+                    raise ToolException(f"Invalid chunking configuration: {e}")
+            #
+            docs = base_chunker(file_content_generator=docs, config=base_chunking_config)
+        #
+        collection_suffix = kwargs.get("collection_suffix")
+        progress_step = kwargs.get("progress_step")
+        clean_index = kwargs.get("clean_index")
+        vs = self._init_vector_store(collection_suffix, embeddings=embedding)
+        #
+        return vs.index_documents(docs, progress_step=progress_step, clean_index=clean_index)
+
     def _process_documents(self, documents: List[Document]) -> Generator[Document, None, None]:
         """
         Process a list of base documents to extract relevant metadata for full document preparation.
@@ -222,10 +304,21 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
             Generator[Document, None, None]: A generator yielding processed documents with metadata.
         """
         for doc in documents:
-            processed_docs = self._process_document(doc)
-            if processed_docs:  # Only proceed if the list is not empty
-                for processed_doc in processed_docs:
-                    yield processed_doc
+            # Filter documents to process only those that either:
+            # - do not have a 'chunk_id' in their metadata, or
+            # - have 'chunk_id' explicitly set to 1.
+            # This prevents processing of irrelevant or duplicate chunks, improving efficiency.
+            chunk_id = doc.metadata.get("chunk_id")
+            if chunk_id is None or chunk_id == 1:
+                processed_docs = self._process_document(doc)
+                if processed_docs:  # Only proceed if the list is not empty
+                    for processed_doc in processed_docs:
+                        # map processed document (child) to the original document (parent)
+                        processed_doc.metadata[IndexerKeywords.PARENT.value] = doc.metadata.get('id', None)
+                        if chunker:=self._get_dependencies_chunker(processed_doc):
+                            yield from chunker(file_content_generator=iter([processed_doc]), config=self._get_dependencies_chunker_config())
+                        else:
+                            yield processed_doc
 
 
     def _init_vector_store(self, collection_suffix: str = "", embeddings: Optional[Any] = None):
@@ -382,6 +475,9 @@ class BaseCodeToolApiWrapper(BaseVectorStoreToolApiWrapper):
     def _read_file(self, file_path: str, branch: str):
         raise NotImplementedError("Subclasses should implement this method")
 
+    def _file_commit_hash(self, file_path: str, branch: str):
+        pass
+
     def __handle_get_files(self, path: str, branch: str):
         """
         Handles the retrieval of files from a specific path and branch.
@@ -447,7 +543,8 @@ class BaseCodeToolApiWrapper(BaseVectorStoreToolApiWrapper):
             for file in _files:
                 if is_whitelisted(file) and not is_blacklisted(file):
                     yield {"file_name": file,
-                           "file_content": self._read_file(file, branch=branch or self.active_branch)}
+                           "file_content": self._read_file(file, branch=branch or self.active_branch),
+                           "commit_hash": self._file_commit_hash(file, branch=branch or self.active_branch)}
 
         return parse_code_files_for_db(file_content_generator())
     
@@ -467,7 +564,7 @@ class BaseCodeToolApiWrapper(BaseVectorStoreToolApiWrapper):
             blacklist=blacklist
         )
         vectorstore = self._init_vector_store(collection_suffix)
-        return vectorstore.index_documents(documents)
+        return vectorstore.index_documents(documents, clean_index=False, is_code=True)
 
     def _get_vector_search_tools(self):
         """
@@ -494,7 +591,11 @@ class BaseCodeToolApiWrapper(BaseVectorStoreToolApiWrapper):
 def extend_with_vector_tools(method):
     def wrapper(self, *args, **kwargs):
         tools = method(self, *args, **kwargs)
-        tools.extend(self._get_vector_search_tools())
+        tools.extend(self._get_vector_search_tools())        
+        #
+        if isinstance(self, BaseVectorStoreToolApiWrapper):
+            tools.append(self.get_index_data_tool())
+        #
         return tools
 
     return wrapper

alita_sdk/tools/figma/__init__.py

@@ -18,7 +18,15 @@ def get_tools(tool):
             oauth2=tool["settings"].get("oauth2", None),
             global_limit=tool["settings"].get("global_limit", GLOBAL_LIMIT),
             global_regexp=tool["settings"].get("global_regexp", None),
-            toolkit_name=tool.get('toolkit_name')
+            toolkit_name=tool.get('toolkit_name'),
+            # indexer settings
+            llm=tool['settings'].get('llm', None),
+            connection_string = tool['settings'].get('connection_string', None),
+            collection_name=str(tool['id']),
+            doctype='doc',
+            embedding_model="HuggingFaceEmbeddings",
+            embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
+            vectorstore_type="PGVector"
         )
         .get_tools()
     )
@@ -45,6 +53,10 @@ class FigmaToolkit(BaseToolkit):
                 List[Literal[tuple(selected_tools)]],
                 Field(default=[], json_schema_extra={"args_schemas": selected_tools}),
             ),
+            # indexer settings
+            connection_string = (Optional[SecretStr], Field(description="Connection string for vectorstore",
+                                                            default=None,
+                                                            json_schema_extra={'secret': True})),
             __config__=ConfigDict(
                 json_schema_extra={
                      "metadata": {

alita_sdk/tools/figma/api_wrapper.py

@@ -1,16 +1,18 @@
+import base64
 import functools
 import json
 import logging
 import re
 from enum import Enum
-from typing import Dict, Optional, Union
+from typing import Dict, Generator, Optional, Union
 
 import requests
 from FigmaPy import FigmaPy
+from langchain_core.documents import Document
 from langchain_core.tools import ToolException
 from pydantic import Field, PrivateAttr, create_model, model_validator, SecretStr
 
-from ..elitea_base import BaseToolApiWrapper
+from ..elitea_base import BaseVectorStoreToolApiWrapper, extend_with_vector_tools
 
 GLOBAL_LIMIT = 10000
 
@@ -226,13 +228,54 @@ class ArgsSchema(Enum):
     )
 
 
-class FigmaApiWrapper(BaseToolApiWrapper):
+class FigmaApiWrapper(BaseVectorStoreToolApiWrapper):
     token: Optional[SecretStr] = Field(default=None)
     oauth2: Optional[SecretStr] = Field(default=None)
     global_limit: Optional[int] = Field(default=GLOBAL_LIMIT)
     global_regexp: Optional[str] = Field(default=None)
     _client: Optional[FigmaPy] = PrivateAttr()
 
+    def _base_loader(self, project_id: str, **kwargs) -> Generator[Document, None, None]:
+        files = json.loads(self.get_project_files(project_id)).get('files', [])
+        for file in files:
+            yield Document(page_content=json.dumps(file), metadata={
+                'id': file.get('key', ''),
+                'file_key': file.get('key', ''),
+                'name': file.get('name', ''),
+                'updated_on': file.get('last_modified', '')
+            })
+
+    def _process_document(self, document: Document) -> Generator[Document, None, None]:
+        file_key = document.metadata.get('id', '')
+        #
+        node_ids = []
+        children = self._client.get_file(file_key).document.get('children', [])
+        if children:
+            nodes = children[0].get('children', [])
+            node_ids = [node['id'] for node in nodes if 'id' in node]
+        images = self._client.get_file_images(file_key, node_ids).images or {}
+
+        # iterate over images values
+        for node_id, image_url in images.items():
+            response = requests.get(image_url)
+            if response.status_code == 200:
+                content_type = response.headers.get('Content-Type', '')
+                if 'text/html' not in content_type.lower():
+                    yield Document(
+                                    page_content=base64.b64encode(response.content).decode("utf-8"),
+                                    metadata={
+                                        'file_key': file_key,
+                                        'node_id': node_id,
+                                        'image_url': image_url
+                                    }
+                                )
+
+    def _index_tool_params(self):
+        """Return the parameters for indexing data."""
+        return {
+            "project_id": (str, Field(description="ID of the project to list files from", examples=["55391681"]))
+        }
+
     def _send_request(
         self,
         method: str,
@@ -442,6 +485,7 @@ class FigmaApiWrapper(BaseToolApiWrapper):
         """Retrieves all files for a specified project ID from Figma."""
         return self._client.get_project_files(project_id)
 
+    @extend_with_vector_tools
     def get_available_tools(self):
         return [
             {

alita_sdk/tools/github/api_wrapper.py

@@ -167,6 +167,14 @@ class AlitaGitHubAPIWrapper(BaseCodeToolApiWrapper):
         # Use the GitHub client's method to get files
         return self.github_client_instance._get_files(path, branch or self.active_branch)
 
+    def _file_commit_hash(self, file_path: str, branch: str):
+        """Get the commit hash of a file in the GitHub repository."""
+        if not self.github_client_instance:
+            raise ValueError("GitHub client not initialized")
+
+        # Use the GitHub client's method to get commit hash
+        return self.github_client_instance._file_commit_hash(file_path, branch or self.active_branch)
+
     def _read_file(self, file_path: str, branch: str):
         """Read file content from GitHub repository."""
         if not self.github_client_instance:

alita_sdk/tools/github/github_client.py

@@ -1393,6 +1393,24 @@ class GitHubClient(BaseModel):
         except Exception as e:
             return f"An error occurred while updating the issue: {str(e)}"
 
+    def _file_commit_hash(self, file_path: str, branch: str, repo_name: Optional[str] = None) -> str:
+        """
+        Get the commit hash of a file in a specific branch.
+        Parameters:
+            file_path(str): the file path
+            branch(str): the branch to read the file from
+            repo_name (Optional[str]): Name of the repository in format 'owner/repo'
+
+        Returns:
+            str: The commit hash of the file, or an error message if not found
+        """
+        try:
+            repo = self.github_api.get_repo(repo_name) if repo_name else self.github_repo_instance
+            file = repo.get_contents(file_path, ref=branch)
+            return file.sha
+        except Exception as e:
+            return f"File not found `{file_path}` on branch `{branch}`. Error: {str(e)}"
+
     def _read_file(self, file_path: str, branch: str, repo_name: Optional[str] = None) -> str:
         """
         Read a file from specified branch

alita_sdk/tools/gitlab/__init__.py

@@ -47,6 +47,10 @@ class AlitaGitlabToolkit(BaseToolkit):
             repository=(str, Field(description="GitLab repository", json_schema_extra={'toolkit_name': True, 'max_toolkit_length': AlitaGitlabToolkit.toolkit_max_length})),
             private_token=(SecretStr, Field(description="GitLab private token", json_schema_extra={'secret': True, 'configuration': True})),
             branch=(str, Field(description="Main branch", default="main")),
+            # indexer settings
+            connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
+                                                          default=None,
+                                                          json_schema_extra={'secret': True})),
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __config__=ConfigDict(json_schema_extra={
                 'metadata': {

alita_sdk/tools/gitlab/api_wrapper.py

@@ -167,6 +167,16 @@ class GitLabAPIWrapper(BaseCodeToolApiWrapper):
         gitlab_files = self._get_all_files(path, recursive, branch)
         return [file['path'] for file in gitlab_files if file['type'] == 'blob']
 
+    def _file_commit_hash(self, file_path: str, branch: str):
+        """
+        Get the commit hash of a file in a specific branch.
+        """
+        try:
+            file = self._repo_instance.files.get(file_path, branch)
+            return file.commit_id
+        except Exception as e:
+            return f"Unable to get commit hash for {file_path} due to error:\n{e}"
+
     def _read_file(self, file_path: str, branch: str):
         return self.read_file(file_path, branch)
 

alita_sdk/tools/google/bigquery/__init__.py

@@ -63,7 +63,7 @@ class BigQueryToolkitConfig(BaseModel):
     dataset: Optional[str] = Field(
         default=None,
         description="BigQuery dataset name",
-        json_schema_extra={"configuration": True},
+        json_schema_extra={"configuration": True, "configuration_title": True},
     )
     table: Optional[str] = Field(
         default=None,

alita_sdk/tools/jira/__init__.py

@@ -11,19 +11,27 @@ name = "jira"
 
 def get_tools(tool):
     return JiraToolkit().get_toolkit(
-            selected_tools=tool['settings'].get('selected_tools', []),
-            base_url=tool['settings'].get('base_url'),
-            cloud=tool['settings'].get('cloud', True),
-            api_key=tool['settings'].get('api_key', None),
-            username=tool['settings'].get('username', None),
-            token=tool['settings'].get('token', None),
-            limit=tool['settings'].get('limit', 5),
-            labels=parse_list(tool['settings'].get('labels', [])),
-            additional_fields=tool['settings'].get('additional_fields', []),
-            toolkit_name=tool.get('toolkit_name'),
-            verify_ssl=tool['settings'].get('verify_ssl', True),
-            llm=tool['settings'].get('llm', None)
-            ).get_tools()
+        selected_tools=tool['settings'].get('selected_tools', []),
+        base_url=tool['settings'].get('base_url'),
+        cloud=tool['settings'].get('cloud', True),
+        api_key=tool['settings'].get('api_key', None),
+        username=tool['settings'].get('username', None),
+        token=tool['settings'].get('token', None),
+        limit=tool['settings'].get('limit', 5),
+        labels=parse_list(tool['settings'].get('labels', [])),
+        additional_fields=tool['settings'].get('additional_fields', []),
+        verify_ssl=tool['settings'].get('verify_ssl', True),
+        # indexer settings
+        llm=tool['settings'].get('llm', None),
+        alita=tool['settings'].get('alita', None),
+        connection_string=tool['settings'].get('connection_string', None),
+        collection_name=f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        doctype='code',
+        embedding_model="HuggingFaceEmbeddings",
+        embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
+        vectorstore_type="PGVector",
+        toolkit_name=tool.get('toolkit_name')
+    ).get_tools()
             
 
 class JiraToolkit(BaseToolkit):