alita-sdk 0.3.114__tar.gz → 0.3.116__tar.gz

{alita_sdk-0.3.114 → alita_sdk-0.3.116}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.114
+Version: 0.3.116
 Summary: SDK for building langchain agents using resouces from Alita
 Author-email: Artem Rozumenko <artyom.rozumenko@gmail.com>, Mikalai Biazruchka <mikalai_biazruchka@epam.com>, Roman Mitusov <roman_mitusov@epam.com>, Ivan Krakhmaliuk <lifedjik@gmail.com>
 Project-URL: Homepage, https://projectalita.ai

{alita_sdk-0.3.114 → alita_sdk-0.3.116}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "alita_sdk"
-version = "0.3.114"
+version = "0.3.116"
 description = "SDK for building langchain agents using resouces from Alita"
 readme = "README.md"
 requires-python = ">=3.10"

{alita_sdk-0.3.114 → alita_sdk-0.3.116}/src/alita_sdk/community/analysis/jira_analyse/__init__.py

@@ -5,6 +5,7 @@ from pydantic import create_model, BaseModel, ConfigDict, Field
 from langchain_core.tools import BaseTool, BaseToolkit
 
 from elitea_analyse.jira.jira_connect import connect_to_jira
+from alita_tools.utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
 from alita_tools.base.tool import BaseAction
 
 from ....tools.artifact import ArtifactWrapper
@@ -16,24 +17,37 @@ name = "Analyse_Jira"
 
 class AnalyseJira(BaseToolkit):
     tools: List[BaseTool] = []
+    toolkit_max_length: int = 0
 
     @staticmethod
     def toolkit_config_schema() -> type[BaseModel]:
+        selected_tools = {x['name']: x['args_schema'].schema() for x in
+                          JiraAnalyseWrapper.model_construct().get_available_tools()}
+        AnalyseJira.toolkit_max_length = get_max_toolkit_length(selected_tools)
         return create_model(
             "analyse_jira",
-            jira_base_url=(str, Field(description="Jira URL")),
+            jira_base_url=(str, Field(
+                description="Jira URL",
+                json_schema_extra={
+                    'toolkit_name': True,
+                    'max_toolkit_length': AnalyseJira.toolkit_max_length                   
+                })
+            ),
             jira_cloud=(bool, Field(description="Hosting Option")),
             jira_username=(str, Field(description="Jira Username")),
             jira_api_key=(Optional[str], Field(description="API key", json_schema_extra={'secret': True}, default="")),
             jira_token=(Optional[str], Field(description="Jira token", json_schema_extra={'secret': True}, default="")),
             # TODO: Add these fields to the schema as custom fields comma-separated if required
+            project_keys=(Optional[str], Field(description="Jira project keys separated by comma", default=None)),
             team_field=(Optional[str], Field(description="Jira field used as identifier for team", default="")),
             environment_field=(Optional[str], Field(description="Jira field used as identifier for environment", default="")),
             defects_name=(Optional[str], Field(description="Jira defects type", default="")),
             closed_status=(Optional[str], Field(description="Jira closed status", default="")),
             jira_verify_ssl=(bool, Field(description="Verify SSL")),
-            jira_custom_fields=(Optional[str], Field(description="Additional fields, split by comma", default="")),
+            jira_custom_fields=(Optional[dict], Field(description="Additional fields, split by comma", default={})),
             artifact_bucket_path=(Optional[str], Field(description="Artifact Bucket Path", default="")),
+            selected_tools=(List[Literal[tuple(selected_tools)]],
+                           Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __config__=ConfigDict(json_schema_extra={'metadata':
                 {
                     "label": "Analyse_Jira",
@@ -59,22 +73,25 @@ class AnalyseJira(BaseToolkit):
         )
 
     @classmethod
-    def get_toolkit(cls, client: 'AlitaClient', **kwargs):
+    def get_toolkit(cls, client: 'AlitaClient', selected_tools: list[str], **kwargs):
+        if selected_tools is None:
+            selected_tools = []
+
         bucket_path = kwargs.get('artifact_bucket_path') or 'analyse-jira'
         artifact_wrapper = ArtifactWrapper(
             client=client, bucket=bucket_path
         )
         check_schema(artifact_wrapper)
 
+        project_keys = kwargs.get('project_keys') or ''
+
         jira_base_url = kwargs.get('jira_base_url')
         jira_verify_ssl = kwargs.get('jira_verify_ssl')
         jira_username = kwargs.get('jira_username')
         jira_token = kwargs.get('jira_token')
         jira_api_key = kwargs.get('jira_api_key')
-        try:
-            jira_custom_fields = json.loads(kwargs.get('jira_custom_fields', '{}'))
-        except:
-            jira_custom_fields = {}
+       
+        jira_custom_fields = kwargs.get('jira_custom_fields', {})
         jira_custom_fields['team'] = kwargs.get('team_field', '')
         jira_custom_fields['environment'] = kwargs.get('environment_field', '')
         closed_status = kwargs.get('closed_status', '')
@@ -97,13 +114,18 @@ class AnalyseJira(BaseToolkit):
         api_wrapper = JiraAnalyseWrapper(
             artifacts_wrapper=artifact_wrapper,
             jira=jira,
+            project_keys=project_keys,
             closed_status=closed_status,
             defects_name=defects_name,
             custom_fields=jira_custom_fields,
         )
+
         tools = []
         available_tools = api_wrapper.get_available_tools()
         for tool in available_tools:
+            if selected_tools:
+                if tool["name"] not in selected_tools:
+                    continue
             tools.append(
                 BaseAction(
                     api_wrapper=api_wrapper,

{alita_sdk-0.3.114 → alita_sdk-0.3.116}/src/alita_sdk/community/analysis/jira_analyse/api_wrapper.py

@@ -10,7 +10,6 @@ import pandas as pd
 
 from elitea_analyse.utils.constants import OUTPUT_MAPPING_FILE, OUTPUT_WORK_ITEMS_FILE
 from elitea_analyse.jira.jira_projects_overview import jira_projects_overview
-from elitea_analyse.jira.jira_all_fields_overview import jira_all_fields_overview
 from elitea_analyse.jira.jira_statuses import get_all_statuses_list
 from elitea_analyse.jira.jira_issues import JiraIssues
 
@@ -22,15 +21,16 @@ logger = logging.getLogger(__name__)
 
 
 class GetJiraFieldsArgs(BaseModel):
-    project_keys: str = Field(
-        description="One or more projects keys separated with comma."
+    project_keys: Optional[str] = Field(
+        description="One or more projects keys separated with comma.",
+        default=''
     )
     after_date: str = Field(description="Date after which issues are considered.")
 
 
 class GetJiraIssuesArgs(BaseModel):
-    project_keys: str = Field(
-        description="One or more projects keys separated with comma."
+    project_keys: Optional[str] = Field(
+        description="One or more projects keys separated with comma.", default=''
     )
     closed_issues_based_on: int = Field(
         description="Define whether issues can be thought as closed based on their status (1) or not empty resolved date (2)."
@@ -46,6 +46,7 @@ class GetJiraIssuesArgs(BaseModel):
 class JiraAnalyseWrapper(BaseToolApiWrapper):
     artifacts_wrapper: ArtifactWrapper
     jira: JIRA
+    project_keys: str  # Jira project keys
     closed_status: str  # Jira ticket closed statuses
     defects_name: str  # Jira ticket defects name
     custom_fields: dict  # Jira ticket custom fields
@@ -53,12 +54,14 @@ class JiraAnalyseWrapper(BaseToolApiWrapper):
     class Config:
         arbitrary_types_allowed = True
 
-    def get_number_off_all_issues(self, project_keys: str, after_date: str):
+    def get_number_off_all_issues(self, after_date: str, project_keys: Optional[str] = None):
         """
         Get projects a user has access to and merge them with issues count.
         after_date: str
             date after which issues are considered
         """
+        project_keys = project_keys or self.project_keys
+
         project_df = jira_projects_overview(
             after_date, project_keys=project_keys, jira=self.jira
         )
@@ -85,69 +88,17 @@ class JiraAnalyseWrapper(BaseToolApiWrapper):
             "projects_summary": project_df.to_string(),
         }
 
-    def get_all_jira_fields(self, project_keys: str, after_date: str):
-        """
-        Get all Jira fields for the specified projects.
-        projects: str
-            one or more projects keys separated with comma
-        after_date: str
-            date after which issues are considered
-        """
-        dispatch_custom_event(
-            name="jira_all_fields_overview",
-            data={
-                "project_keys": project_keys,
-                "after_date": after_date,
-            },
-        )
-        overall_stat, issue_types_stat = jira_all_fields_overview(
-            project_keys, after_date, jira=self.jira
-        )
-
-        dispatch_custom_event(
-            name="jira_fields_saving",
-            data={
-                "project_keys": project_keys,
-                "after_date": after_date,
-                "overall_stat": len(overall_stat),
-                "issue_types_stat": len(issue_types_stat),
-                "files": [
-                    "fields_count.csv",
-                    f"fields_count_issues_{project_keys}.csv",
-                ],
-            },
-        )
-
-        self.save_dataframe(
-            overall_stat,
-            "fields_count.csv",
-            csv_options={"index": False},
-        )
-
-        self.save_dataframe(
-            issue_types_stat,
-            f"fields_count_issues_{project_keys}.csv",
-            csv_options={"index": False},
-        )
-
-        return {
-            "overall_stat": overall_stat.to_string(),
-            "issue_types_stat": issue_types_stat.to_string(),
-        }
-
     def get_jira_issues(
         self,
-        project_keys: str,
         closed_issues_based_on: int,
         resolved_after: str,
         updated_after: str,
         created_after: str,
         add_filter: str = "",
+        project_keys: Optional[str] = None,
     ):
         """
         Extract Jira issues for the specified projects.
-        projects: str
-            one or more projects keys separated with comma
         closed_issues_based_on: int
             define whether issues can be thought as closed based on their status (1) or not empty resolved date (2)
         resolved_after: str
@@ -158,8 +109,9 @@ class JiraAnalyseWrapper(BaseToolApiWrapper):
             created after date (i.e. 2023-01-01)
         add_filter: str
             additional filter for Jira issues in JQL format like "customfield_10000 = 'value' AND customfield_10001 = 'value'"
+        project_keys: str
+            one or more projects keys separated with comma
         """
-
         if not (
             (
                 closed_issues_based_on == 1
@@ -170,14 +122,22 @@ class JiraAnalyseWrapper(BaseToolApiWrapper):
             return (
                 "ERROR: Check input parameters closed_issues_based_on and closed_status"
             )
+        
+        project_keys = project_keys or self.project_keys
 
         dispatch_custom_event(
-            name="jira_issues_extraction_start",
+            name="thinking_step",
             data={
-                "closed_issues_based_on": closed_issues_based_on,
-                "closed_status": self.closed_status,
-            }
+                "message": f"I am extracting Jira issues with initial parameters:\
+                    project keys: {project_keys}, closed status: {self.closed_status},\
+                    defects name: {self.defects_name}, custom fields: {self.custom_fields}, \
+                    closed status based on: {closed_issues_based_on}, resolved after: {resolved_after}, \
+                    updated after: {updated_after}, created after: {created_after}, additional filter:{add_filter}",
+                "tool_name": "jira_issues_extraction_start",
+                "toolkit": "analyse_jira",
+            },
         )
+
         jira_issues = JiraIssues(
             self.jira,
             project_keys,
@@ -189,12 +149,15 @@ class JiraAnalyseWrapper(BaseToolApiWrapper):
         df_issues, df_map = jira_issues.extract_issues_from_jira_and_transform(
             self.custom_fields, (resolved_after, updated_after, created_after)
         )
+
         dispatch_custom_event(
-            name="jira_issues_extracted",
+            name="thinking_step",
             data={
-                "project_keys": jira_issues.projects,
-                "issue_count": len(df_issues),
-                "map_rows": len(df_map),
+                "message": f"I am saving the extracted Jira issues to the artifact repository. \
+                            issues count: {len(df_issues)}, mapping rows: {len(df_map)}, \
+                            output file: {OUTPUT_MAPPING_FILE}{jira_issues.projects}.csv",
+                "tool_name": "get_jira_issues",
+                "toolkit": "analyse_jira",
             },
         )
         self.save_dataframe(
@@ -202,13 +165,6 @@ class JiraAnalyseWrapper(BaseToolApiWrapper):
             f"{OUTPUT_MAPPING_FILE}{jira_issues.projects}.csv",
             csv_options={"index_label": "id"},
         )
-        dispatch_custom_event(
-            name="jira_map_statuces_saved",
-            data={
-                "output_file": f"{OUTPUT_MAPPING_FILE}{jira_issues.projects}.csv",
-                "row_count": len(df_map),
-            }
-        )
 
         if not df_issues.empty:
             self.save_dataframe(
@@ -217,11 +173,14 @@ class JiraAnalyseWrapper(BaseToolApiWrapper):
                 csv_options={"index_label": "id"},
             )
             dispatch_custom_event(
-                name="jira_issues_saved",
+                name="thinking_step",
                 data={
-                    "output_file": f"{OUTPUT_WORK_ITEMS_FILE}{jira_issues.projects}.csv",
-                    "row_count": len(df_issues),
-                }
+                    "message": f"Saving Jira issues to the file . \
+                                 output file: {OUTPUT_WORK_ITEMS_FILE}{jira_issues.projects}.csv,\
+                                 row count: {len(df_issues)}",
+                    "tool_name": "get_jira_issues",
+                    "toolkit": "analyse_jira",
+                },
             )
 
         return f"{jira_issues.projects} Data has been extracted successfully."
@@ -235,12 +194,6 @@ class JiraAnalyseWrapper(BaseToolApiWrapper):
                 "args_schema": GetJiraFieldsArgs,
                 "ref": self.get_number_off_all_issues,
             },
-            {
-                "name": "get_all_jira_fields",
-                "description": self.get_all_jira_fields.__doc__,
-                "args_schema": GetJiraFieldsArgs,
-                "ref": self.get_all_jira_fields,
-            },
             {
                 "name": "get_jira_issues",
                 "description": self.get_jira_issues.__doc__,

{alita_sdk-0.3.114 → alita_sdk-0.3.116}/src/alita_sdk/langchain/langraph_agent.py

@@ -26,6 +26,7 @@ from ..tools.loop_output import LoopToolNode
 from ..tools.tool import ToolNode
 from ..utils.evaluate import EvaluateTemplate
 from ..utils.utils import clean_string, TOOLKIT_SPLITTER
+from ..tools.router import RouterNode
 
 logger = logging.getLogger(__name__)
 
@@ -124,6 +125,58 @@ class TransitionalEdge(Runnable):
         return self.next_step if self.next_step != 'END' else END
 
 
+class StateModifierNode(Runnable):
+    name = "StateModifierNode"
+
+    def __init__(self, template: str, variables_to_clean: Optional[list[str]] = None, 
+                 input_variables: Optional[list[str]] = None, 
+                 output_variables: Optional[list[str]] = None):
+        self.template = template
+        self.variables_to_clean = variables_to_clean or []
+        self.input_variables = input_variables or ["messages"]
+        self.output_variables = output_variables or []
+
+    def invoke(self, state: Annotated[BaseStore, InjectedStore()], config: Optional[RunnableConfig] = None) -> dict:
+        logger.info(f"Modifying state with template: {self.template}")
+
+        # Collect input variables from state
+        input_data = {}
+        for var in self.input_variables:
+            if var in state:
+                input_data[var] = state.get(var)
+
+        # Render the template using Jinja
+        from jinja2 import Template
+        rendered_message = Template(self.template).render(**input_data)
+        result = {}
+        # Store the rendered message in the state or messages
+        if len(self.output_variables) > 0:
+            # Use the first output variable to store the rendered content
+            output_var = self.output_variables[0]
+            result[output_var] = rendered_message
+
+        # Clean up specified variables (make them empty, not delete)
+        
+        for var in self.variables_to_clean:
+            if var in state:
+                # Empty the variable based on its type
+                if isinstance(state[var], list):
+                    result[var] = []
+                elif isinstance(state[var], dict):
+                    result[var] = {}
+                elif isinstance(state[var], str):
+                    result[var] = ""
+                elif isinstance(state[var], (int, float)):
+                    result[var] = 0
+                elif state[var] is None:
+                    pass
+                else:
+                    # For other types, set to None
+                    result[var] = None
+        logger.info(f"State modifier result: {result}")
+        return result
+
+
 def prepare_output_schema(lg_builder, memory, store, debug=False, interrupt_before=[], interrupt_after=[]):
     # prepare output channels
     output_channels = (
@@ -283,6 +336,32 @@ def create_graph(
                     output_variables=node.get('output', []),
                     input_variables=node.get('input', ['messages']),
                     structured_output=node.get('structured_output', False)))
+            elif node_type == 'router':
+                # Add a RouterNode as an independent node
+                lg_builder.add_node(node_id, RouterNode(
+                    name=node['id'],
+                    condition=node.get('condition', ''),
+                    routes=node.get('routes', []),
+                    default_output=node.get('default_output', 'END'),
+                    input_variables=node.get('input', ['messages'])
+                ))
+                # Add a single conditional edge for all routes
+                lg_builder.add_conditional_edges(
+                    node_id,
+                    ConditionalEdge(
+                        condition="{{router_output}}",  # router node returns the route key in 'router_output'
+                        condition_inputs=["router_output"],
+                        conditional_outputs=node.get('routes', []),
+                        default_output=node.get('default_output', 'END')
+                    )
+                )
+            elif node_type == 'state_modifier':
+                lg_builder.add_node(node_id, StateModifierNode(
+                    template=node.get('template', ''),
+                    variables_to_clean=node.get('variables_to_clean', []),
+                    input_variables=node.get('input', ['messages']),
+                    output_variables=node.get('output', [])
+                ))
             if node.get('transition'):
                 next_step = clean_string(node['transition'])
                 logger.info(f'Adding transition: {next_step}')

{alita_sdk-0.3.114 → alita_sdk-0.3.116}/src/alita_sdk/langchain/utils.py

@@ -128,7 +128,7 @@ def parse_type(type_str):
 
 
 def create_state(data: Optional[dict] = None):
-    state_dict = {'input': str,}
+    state_dict = {'input': str, 'router_output': str}  # Always include router_output
     if not data:
         data = {'messages': 'list[str]'}
     for key, value in data.items():

alita_sdk-0.3.116/src/alita_sdk/tools/router.py

@@ -0,0 +1,35 @@
+import logging
+from typing import Any, Optional, Union, List
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import BaseTool
+from ..utils.evaluate import EvaluateTemplate
+from ..utils.utils import clean_string
+
+logger = logging.getLogger(__name__)
+
+class RouterNode(BaseTool):
+    name: str = 'RouterNode'
+    description: str = 'A router node that evaluates a condition and routes accordingly.'
+    condition: str = ''
+    routes: List[str] = []  # List of possible output node keys
+    default_output: str = 'END'
+    input_variables: Optional[list[str]] = None
+
+    def invoke(self, state: Union[str, dict], config: Optional[RunnableConfig] = None, **kwargs: Any) -> dict:
+        input_data = {}
+        for field in self.input_variables or []:
+            input_data[field] = state.get(field, "")
+        template = EvaluateTemplate(self.condition, input_data)
+        result = template.evaluate()
+        logger.info(f"RouterNode evaluated condition '{self.condition}' with input {input_data} => {result}")
+        result = clean_string(str(result))
+        if result in self.routes:
+            # If the result is one of the routes, return it
+            return {"router_output": result}
+        elif result == self.default_output:
+            # If the result is the default output, return it
+            return {"router_output": clean_string(self.default_output)}
+        return {"router_output": 'END'}
+
+    def _run(self, *args, **kwargs):
+        return self.invoke(**kwargs)

{alita_sdk-0.3.114 → alita_sdk-0.3.116}/src/alita_sdk/utils/AlitaCallback.py

@@ -1,6 +1,7 @@
 import logging
 import json
 import traceback
+from datetime import datetime, timezone
 from uuid import UUID, uuid4
 from typing import Any, Dict, List, Optional
 from collections import defaultdict
@@ -51,6 +52,35 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
     # Tool
     #
 
+    def on_custom_event(
+        self,
+        name: str,
+        data: Any,
+        *,
+        run_id: UUID,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Callback containing a group of custom events"""
+
+        payload = {
+            "name": name,
+            "run_id": str(run_id),
+            "tool_run_id": str(run_id),  # compatibility
+            "metadata": metadata,
+            "datetime": str(datetime.now(tz=timezone.utc)),
+            **data,
+        }
+        payload = json.loads(
+            json.dumps(payload, ensure_ascii=False, default=lambda o: str(o))
+        )
+
+        self.callback_state[str(run_id)] = self.st.status(
+            f"Running {payload.get("tool_name")}...", expanded=True
+        )
+        self.callback_state[str(run_id)].write(f"Tool inputs: {payload}")
+
     def on_tool_start(self, *args, run_id: UUID, **kwargs):
         """ Callback """
         if self.debug:
@@ -125,7 +155,7 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
 
         self.current_model_name = metadata.get('ls_model_name', self.current_model_name)
         llm_run_id = str(run_id)
-        
+
         self.callback_state[llm_run_id] = self.st.status(f"Running LLM ...", expanded=True)
         self.callback_state[llm_run_id].write(f"LLM inputs: {messages}")
 
@@ -176,4 +206,4 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
             log.debug("on_llm_end(%s, %s)", response, kwargs)
         llm_run_id = str(run_id)
         self.callback_state[llm_run_id].update(label=f"Completed LLM call", state="complete", expanded=False)
-        self.callback_state.pop(llm_run_id, None)
+        self.callback_state.pop(llm_run_id, None)

{alita_sdk-0.3.114 → alita_sdk-0.3.116}/src/alita_sdk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.114
+Version: 0.3.116
 Summary: SDK for building langchain agents using resouces from Alita
 Author-email: Artem Rozumenko <artyom.rozumenko@gmail.com>, Mikalai Biazruchka <mikalai_biazruchka@epam.com>, Roman Mitusov <roman_mitusov@epam.com>, Ivan Krakhmaliuk <lifedjik@gmail.com>
 Project-URL: Homepage, https://projectalita.ai

{alita_sdk-0.3.114 → alita_sdk-0.3.116}/src/alita_sdk.egg-info/SOURCES.txt

@@ -91,6 +91,7 @@ src/alita_sdk/tools/loop.py
 src/alita_sdk/tools/loop_output.py
 src/alita_sdk/tools/pgvector_search.py
 src/alita_sdk/tools/prompt.py
+src/alita_sdk/tools/router.py
 src/alita_sdk/tools/tool.py
 src/alita_sdk/tools/vectorstore.py
 src/alita_sdk/utils/AlitaCallback.py