alita-sdk 0.3.203__py3-none-any.whl → 0.3.205__py3-none-any.whl

alita_sdk/tools/elitea_base.py

@@ -3,6 +3,8 @@ import fnmatch
 import logging
 import traceback
 from typing import Any, Optional, List, Dict
+
+from langchain_core.documents import Document
 from langchain_core.tools import ToolException
 from pydantic import BaseModel, create_model, Field
 from .utils import TOOLKIT_SPLITTER
@@ -186,7 +188,41 @@ class BaseToolApiWrapper(BaseModel):
 
 class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
     """Base class for tool API wrappers that support vector store functionality."""
-    
+
+    doctype: str = "document"
+
+    def _base_loader(self, **kwargs) -> List[Document]:
+        """ 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:
+        """ 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.
+
+        Args:
+            document (Document): The base document to process.
+
+        Returns:
+            Document: The processed document with metadata."""
+        pass
+
+    def _process_documents(self, documents: List[Document]) -> List[Document]:
+        """
+        Process a list of base documents to extract relevant metadata for full document preparation.
+        Used for late processing of documents after we ensure that the documents have to be indexed to avoid
+        time-consuming operations for documents which might be useless.
+        This function passed to index_documents method of vector store and called after _reduce_duplicates method.
+
+        Args:
+            documents (List[Document]): The base documents to process.
+
+        Returns:
+            List[Document]: The processed documents with metadata.
+        """
+        return [self._process_document(doc) for doc in documents]
+
     def _init_vector_store(self, collection_suffix: str = "", embeddings: Optional[Any] = None):
         """ Initializes the vector store wrapper with the provided parameters."""
         try:
@@ -225,7 +261,8 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
             embedding_model=self.embedding_model,
             embedding_model_params=self.embedding_model_params,
             vectorstore_params=vectorstore_params,
-            embeddings=embeddings
+            embeddings=embeddings,
+            process_document_func=self._process_documents,
         )
 
     def search_index(self,
@@ -345,7 +382,7 @@ class BaseCodeToolApiWrapper(BaseVectorStoreToolApiWrapper):
         Handles the retrieval of files from a specific path and branch.
         This method should be implemented in subclasses to provide the actual file retrieval logic.
         """
-        _files = self._get_files(path, branch)
+        _files = self._get_files(path=path, branch=branch)
         if isinstance(_files, str):
             try:
                 # Attempt to convert the string to a list using ast.literal_eval
@@ -447,4 +484,12 @@ class BaseCodeToolApiWrapper(BaseVectorStoreToolApiWrapper):
         }
         
         # Return index_data tool first, then the search tools
-        return [index_tool] + base_tools
+        return [index_tool] + base_tools
+
+def extend_with_vector_tools(method):
+    def wrapper(self, *args, **kwargs):
+        tools = method(self, *args, **kwargs)
+        tools.extend(self._get_vector_search_tools())
+        return tools
+
+    return wrapper

alita_sdk/tools/google/__init__.py

@@ -0,0 +1,7 @@
+from .bigquery import BigQueryToolkit
+
+name = "google"
+
+def get_tools(tool_type, tool):
+    if tool_type == 'bigquery':
+        return BigQueryToolkit().get_toolkit().get_tools()

alita_sdk/tools/google/bigquery/__init__.py

@@ -0,0 +1,154 @@
+from functools import lru_cache
+from typing import List, Optional, Type
+
+from langchain_core.tools import BaseTool, BaseToolkit
+from pydantic import BaseModel, Field, SecretStr, computed_field, field_validator
+
+from ...utils import TOOLKIT_SPLITTER, clean_string, get_max_toolkit_length
+from .api_wrapper import BigQueryApiWrapper
+from .tool import BigQueryAction
+
+name = "bigquery"
+
+
+@lru_cache(maxsize=1)
+def get_available_tools() -> dict[str, dict]:
+    api_wrapper = BigQueryApiWrapper.model_construct()
+    available_tools: dict = {
+        x["name"]: x["args_schema"].model_json_schema()
+        for x in api_wrapper.get_available_tools()
+    }
+    return available_tools
+
+
+toolkit_max_length = lru_cache(maxsize=1)(
+    lambda: get_max_toolkit_length(get_available_tools())
+)
+
+
+class BigQueryToolkitConfig(BaseModel):
+    class Config:
+        title = name
+        json_schema_extra = {
+            "metadata": {
+                "hidden": True,
+                "label": "Cloud GCP",
+                "icon_url": "google.svg",
+                "sections": {
+                    "auth": {
+                        "required": False,
+                        "subsections": [
+                            {"name": "API Key", "fields": ["api_key"]},
+                        ],
+                    }
+                },
+            }
+        }
+
+    api_key: Optional[SecretStr] = Field(
+        default=None,
+        description="GCP API key",
+        json_schema_extra={"secret": True, "configuration": True},
+    )
+    project: Optional[str] = Field(
+        default=None,
+        description="BigQuery project ID",
+        json_schema_extra={"configuration": True},
+    )
+    location: Optional[str] = Field(
+        default=None,
+        description="BigQuery location",
+        json_schema_extra={"configuration": True},
+    )
+    dataset: Optional[str] = Field(
+        default=None,
+        description="BigQuery dataset name",
+        json_schema_extra={"configuration": True},
+    )
+    table: Optional[str] = Field(
+        default=None,
+        description="BigQuery table name",
+        json_schema_extra={"configuration": True},
+    )
+    selected_tools: List[str] = Field(
+        default=[],
+        description="Selected tools",
+        json_schema_extra={"args_schemas": get_available_tools()},
+    )
+
+    @field_validator("selected_tools", mode="before", check_fields=False)
+    @classmethod
+    def selected_tools_validator(cls, value: List[str]) -> list[str]:
+        return [i for i in value if i in get_available_tools()]
+
+
+def _get_toolkit(tool) -> BaseToolkit:
+    return BigQueryToolkit().get_toolkit(
+        selected_tools=tool["settings"].get("selected_tools", []),
+        api_key=tool["settings"].get("api_key", ""),
+        toolkit_name=tool.get("toolkit_name"),
+    )
+
+
+def get_toolkit():
+    return BigQueryToolkit.toolkit_config_schema()
+
+
+def get_tools(tool):
+    return _get_toolkit(tool).get_tools()
+
+
+class BigQueryToolkit(BaseToolkit):
+    tools: List[BaseTool] = []
+    api_wrapper: Optional[BigQueryApiWrapper] = Field(
+        default_factory=BigQueryApiWrapper.model_construct
+    )
+    toolkit_name: Optional[str] = None
+
+    @computed_field
+    @property
+    def tool_prefix(self) -> str:
+        return (
+            clean_string(self.toolkit_name, toolkit_max_length()) + TOOLKIT_SPLITTER
+            if self.toolkit_name
+            else ""
+        )
+
+    @computed_field
+    @property
+    def available_tools(self) -> List[dict]:
+        return self.api_wrapper.get_available_tools()
+
+    @staticmethod
+    def toolkit_config_schema() -> Type[BaseModel]:
+        return BigQueryToolkitConfig
+
+    @classmethod
+    def get_toolkit(
+        cls,
+        selected_tools: list[str] | None = None,
+        toolkit_name: Optional[str] = None,
+        **kwargs,
+    ) -> "BigQueryToolkit":
+        bigquery_api_wrapper = BigQueryApiWrapper(**kwargs)
+        instance = cls(
+            tools=[], api_wrapper=bigquery_api_wrapper, toolkit_name=toolkit_name
+        )
+        if selected_tools:
+            selected_tools = set(selected_tools)
+            for t in instance.available_tools:
+                if t["name"] in selected_tools:
+                    instance.tools.append(
+                        BigQueryAction(
+                            api_wrapper=instance.api_wrapper,
+                            name=instance.tool_prefix + t["name"],
+                            # set unique description for declared tools to differentiate the same methods for different toolkits
+                            description=f"Project: {getattr(instance.api_wrapper, 'project', '')}\n"
+                            + t["description"],
+                            args_schema=t["args_schema"],
+                        )
+                    )
+        return instance
+
+    def get_tools(self):
+        return self.tools