alita-sdk 0.3.148__py3-none-any.whl → 0.3.150__py3-none-any.whl

alita_sdk/runtime/clients/client.py

@@ -9,6 +9,7 @@ from langchain_core.messages import (
     SystemMessage, BaseMessage,
 )
 from langchain_core.tools import ToolException
+from langgraph.store.base import BaseStore
 
 from ..langchain.assistant import Assistant as LangChainAssistant
 # from ..llamaindex.assistant import Assistant as LLamaAssistant
@@ -178,7 +179,7 @@ class AlitaClient:
                     tools: Optional[list] = None, chat_history: Optional[List[Any]] = None,
                     app_type=None, memory=None, runtime='langchain',
                     application_variables: Optional[dict] = None,
-                    version_details: Optional[dict] = None):
+                    version_details: Optional[dict] = None, store: Optional[BaseStore] = None):
         if tools is None:
             tools = []
         if chat_history is None:
@@ -210,11 +211,11 @@ class AlitaClient:
             app_type = "openai"
         if runtime == 'nonrunnable':
             return LangChainAssistant(self, data, client, chat_history, app_type,
-                                      tools=tools, memory=memory)
+                                      tools=tools, memory=memory, store=store)
         if runtime == 'langchain':
             return LangChainAssistant(self, data, client,
                                       chat_history, app_type,
-                                      tools=tools, memory=memory).runnable()
+                                      tools=tools, memory=memory, store=store).runnable()
         elif runtime == 'llama':
             raise NotImplementedError("LLama runtime is not supported")
 

alita_sdk/runtime/langchain/assistant.py

@@ -5,8 +5,9 @@ from typing import Any, Optional
 from langchain.agents import (
     AgentExecutor, create_openai_tools_agent,
     create_json_chat_agent)
+from langgraph.store.base import BaseStore
+
 from .agents.xml_chat import create_xml_chat_agent
-# from ..langchain.mixedAgentRenderes import render_react_text_description_and_args
 from .langraph_agent import create_graph
 from langchain_core.messages import (
     BaseMessage, SystemMessage, HumanMessage
@@ -28,7 +29,8 @@ class Assistant:
                  chat_history: list[BaseMessage] = [],
                  app_type: str = "openai",
                  tools: Optional[list] = [],
-                 memory: Optional[Any] = None):
+                 memory: Optional[Any] = None,
+                 store: Optional[BaseStore] = None):
 
         self.client = copy(client)
         self.client.max_tokens = data['llm_settings']['max_tokens']
@@ -40,6 +42,7 @@ class Assistant:
 
         self.app_type = app_type
         self.memory = memory
+        self.store = store
 
         logger.debug("Data for agent creation: %s", data)
         logger.info("App type: %s", app_type)
@@ -58,7 +61,10 @@ class Assistant:
             raise ToolException("Non-pipeline agents cannot have pipelines as a toolkits. "
                                 "Review toolkits configuration or use pipeline as master agent.")
 
-        self.tools = get_tools(data['tools'], alita_client=alita, llm=self.client)
+        # configure memory store if memory tool is defined
+        memory_tool = next((tool for tool in data['tools'] if tool['type'] == 'memory'), None)
+        self._configure_store(memory_tool)
+        self.tools = get_tools(data['tools'], alita_client=alita, llm=self.client, memory_store=self.store)
         if app_type == "pipeline":
             self.prompt = data['instructions']
         else:
@@ -94,6 +100,18 @@ class Assistant:
             except Exception as e:
                 logger.info(f"Client was created with client setting: temperature - {self.client.temperature} : {self.client.max_tokens}")
 
+    def _configure_store(self, memory_tool: dict | None) -> None:
+        """
+        Configure the memory store based on a memory_tool definition.
+        Only creates a new store if one does not already exist.
+        """
+        if not memory_tool or self.store is not None:
+            return
+        from .store_manager import get_manager
+        conn_str = memory_tool.get('settings', {}).get('connection_string', '')
+        store = get_manager().get_store(conn_str)
+        self.store = store
+
     def runnable(self):
         if self.app_type == 'pipeline':
             return self.pipeline()

alita_sdk/runtime/langchain/store_manager.py

@@ -0,0 +1,46 @@
+import threading
+import atexit
+import logging
+from psycopg import Connection
+from langgraph.store.postgres import PostgresStore
+
+logger = logging.getLogger(__name__)
+
+class StoreManager:
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._stores = {}
+        return cls._instance
+
+    def get_store(self, conn_str: str) -> PostgresStore:
+        store = self._stores.get(conn_str)
+        if store is None:
+            logger.info(f"Creating new PostgresStore for connection: {conn_str}")
+            conn = Connection.connect(conn_str, autocommit=True, prepare_threshold=0)
+            store = PostgresStore(conn)
+            store.setup()
+            self._stores[conn_str] = store
+        return store
+
+    def shutdown(self) -> None:
+        logger.info("Shutting down StoreManager and closing all stores")
+        for store in list(self._stores.values()):
+            try:
+                conn = getattr(store, 'conn', None)
+                if conn:
+                    conn.close()
+            except Exception:
+                pass
+        self._stores.clear()
+
+_store_manager = StoreManager()
+atexit.register(_store_manager.shutdown)
+
+def get_manager() -> StoreManager:
+    return _store_manager

alita_sdk/runtime/toolkits/application.py

@@ -1,4 +1,6 @@
 from typing import List, Any, Optional
+
+from langgraph.store.base import BaseStore
 from pydantic import create_model, BaseModel, Field
 from langchain_community.agent_toolkits.base import BaseToolkit
 from langchain_core.tools import BaseTool
@@ -19,8 +21,8 @@ class ApplicationToolkit(BaseToolkit):
         )
     
     @classmethod
-    def get_toolkit(cls, client: Any, application_id: int, application_version_id: int, app_api_key: str, 
-                    selected_tools: list[str] = []):
+    def get_toolkit(cls, client: Any, application_id: int, application_version_id: int, app_api_key: str,
+                    selected_tools: list[str] = [], store: Optional[BaseStore] = None):
         from ..llms.alita import AlitaChatModel
         
         app_details = client.get_app_details(application_id)
@@ -37,7 +39,7 @@ class ApplicationToolkit(BaseToolkit):
             "temperature": version_details['llm_settings']['temperature'],
         }
 
-        app = client.application(AlitaChatModel(**settings), application_id, application_version_id)
+        app = client.application(AlitaChatModel(**settings), application_id, application_version_id, store=store)
         return cls(tools=[Application(name=app_details.get("name"), 
                                       description=app_details.get("description"), 
                                       application=app, 

alita_sdk/runtime/toolkits/tools.py

@@ -1,19 +1,20 @@
 import logging
 
+from langchain_core.tools import ToolException
+from langgraph.store.base import BaseStore
+
 from alita_sdk.tools import get_toolkits as alita_toolkits
 from alita_sdk.tools import get_tools as alita_tools
-
 from .application import ApplicationToolkit
 from .artifact import ArtifactToolkit
 from .datasource import DatasourcesToolkit
 from .prompt import PromptToolkit
 from .subgraph import SubgraphToolkit
 from .vectorstore import VectorStoreToolkit
-
+from ..tools.mcp_server_tool import McpServerTool
 # Import community tools
 from ...community import get_toolkits as community_toolkits, get_tools as community_tools
-
-from ..tools.mcp_server_tool import McpServerTool
+from ...tools.memory import MemoryToolkit
 
 logger = logging.getLogger(__name__)
 
@@ -24,13 +25,14 @@ def get_toolkits():
         # DatasourcesToolkit.toolkit_config_schema(),
         # ApplicationToolkit.toolkit_config_schema(),
         ArtifactToolkit.toolkit_config_schema(),
+        MemoryToolkit.toolkit_config_schema(),
         VectorStoreToolkit.toolkit_config_schema()
     ]
 
     return core_toolkits + community_toolkits() + alita_toolkits()
 
 
-def get_tools(tools_list: list, alita_client, llm) -> list:
+def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = None) -> list:
     prompts = []
     tools = []
 
@@ -65,6 +67,14 @@ def get_tools(tools_list: list, alita_client, llm) -> list:
                 selected_tools=[],
                 llm=llm
             ))
+        elif tool['type'] == 'memory':
+            if memory_store is None:
+                raise ToolException(f"Memory store is not provided for memory tool: {tool['name']}")
+            tools += MemoryToolkit.get_toolkit(
+                namespace=tool['settings'].get('namespace', str(tool['id'])),
+                store=memory_store,
+                toolkit_name=tool.get('toolkit_name', '')
+            ).get_tools()
         elif tool['type'] == 'artifact':
             tools.extend(ArtifactToolkit.get_toolkit(
                 client=alita_client,

alita_sdk/tools/memory/__init__.py

@@ -0,0 +1,40 @@
+from typing import Optional, List
+
+from langchain_core.tools import BaseToolkit, BaseTool
+from langgraph.store.postgres import PostgresStore
+from langmem import create_manage_memory_tool, create_search_memory_tool
+from pydantic import create_model, BaseModel, ConfigDict, Field, SecretStr
+
+name = "memory"
+
+class MemoryToolkit(BaseToolkit):
+    tools: List[BaseTool] = []
+
+
+    @staticmethod
+    def toolkit_config_schema() -> BaseModel:
+        return create_model(
+            name,
+            namespace=(str, Field(description="Memory namespace", json_schema_extra={'toolkit_name': True})),
+            connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
+                                                          default=None,
+                                                          json_schema_extra={'secret': True})),
+            __config__=ConfigDict(json_schema_extra={
+                'metadata': {
+                    "label": "Memory",
+                    "icon_url": "jira-icon.svg",
+                    "categories": ["other"],
+                    "extra_categories": ["long-term memory", "langmem"],
+                }
+            })
+        )
+
+    @classmethod
+    def get_toolkit(cls, namespace: str, store: PostgresStore, **kwargs):
+        return cls(tools=[
+            create_manage_memory_tool(namespace=namespace, store=store),
+            create_search_memory_tool(namespace=namespace, store=store)
+        ])
+
+    def get_tools(self):
+        return self.tools

{alita_sdk-0.3.148.dist-info → alita_sdk-0.3.150.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.148
+Version: 0.3.150
 Summary: SDK for building langchain agents using resources 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>
 License-Expression: Apache-2.0
@@ -52,7 +52,7 @@ Requires-Dist: openpyxl==3.1.2; extra == "runtime"
 Requires-Dist: markdown==3.5.1; extra == "runtime"
 Requires-Dist: beautifulsoup4==4.12.2; extra == "runtime"
 Requires-Dist: gensim==4.3.3; extra == "runtime"
-Requires-Dist: chromadb<0.7.0; extra == "runtime"
+Requires-Dist: chromadb==0.5.20; extra == "runtime"
 Requires-Dist: pgvector==0.2.5; extra == "runtime"
 Requires-Dist: scipy==1.13.1; extra == "runtime"
 Requires-Dist: pytesseract==0.3.13; extra == "runtime"
@@ -62,7 +62,11 @@ Requires-Dist: rlpycairo==0.3.0; extra == "runtime"
 Requires-Dist: cairocffi==1.7.1; extra == "runtime"
 Requires-Dist: docx2txt==0.8; extra == "runtime"
 Requires-Dist: mammoth==1.9.0; extra == "runtime"
-Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>1.11.1; extra == "runtime"
+Requires-Dist: opentelemetry-exporter-otlp-proto-grpc==1.25.0; extra == "runtime"
+Requires-Dist: opentelemetry_api==1.25.0; extra == "runtime"
+Requires-Dist: opentelemetry_instrumentation==0.46b0; extra == "runtime"
+Requires-Dist: grpcio_status==1.63.0rc1; extra == "runtime"
+Requires-Dist: protobuf==4.25.7; extra == "runtime"
 Requires-Dist: sentence-transformers==2.7.0; extra == "runtime"
 Provides-Extra: tools
 Requires-Dist: dulwich==0.21.6; extra == "tools"

{alita_sdk-0.3.148.dist-info → alita_sdk-0.3.150.dist-info}/RECORD

@@ -44,17 +44,18 @@ alita_sdk/community/deep_researcher/utils/os.py,sha256=Q1xX7c7_p7EmuzzXIAY9TDmra
 alita_sdk/runtime/__init__.py,sha256=-2xVy-AmrOgiqyA8atCeiDo85GPtEV0robUSL7XIplc,796
 alita_sdk/runtime/clients/__init__.py,sha256=5O_7WsZmvg5z5uZf_Jkx-f8j5C6yKIuSwQR6AQ-TM84,31
 alita_sdk/runtime/clients/artifact.py,sha256=33prjst8z3Wn3SZ8Xl2gJIgQmKzaJPDyaVzay87mDes,2643
-alita_sdk/runtime/clients/client.py,sha256=1QZp6Is1WA0z1EPdqQY6t-_rt7VBBsHnCSRznM6teyY,19630
+alita_sdk/runtime/clients/client.py,sha256=jbC_M72CybwZgFfMRL6paj-NmICrSuk1vVnVTm_u-kc,19734
 alita_sdk/runtime/clients/datasource.py,sha256=HAZovoQN9jBg0_-lIlGBQzb4FJdczPhkHehAiVG3Wx0,1020
 alita_sdk/runtime/clients/prompt.py,sha256=li1RG9eBwgNK_Qf0qUaZ8QNTmsncFrAL2pv3kbxZRZg,1447
 alita_sdk/runtime/langchain/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-alita_sdk/runtime/langchain/assistant.py,sha256=Wq4_XtlFKKkrkelvTa4qzkI3BBTHy5zwRmk3v1i22lQ,6676
+alita_sdk/runtime/langchain/assistant.py,sha256=DjMgup0nfFe5jgNkC-xGgZPiAg7wiJ2qb3VN06n_fnE,7439
 alita_sdk/runtime/langchain/chat_message_template.py,sha256=kPz8W2BG6IMyITFDA5oeb5BxVRkHEVZhuiGl4MBZKdc,2176
 alita_sdk/runtime/langchain/constants.py,sha256=eHVJ_beJNTf1WJo4yq7KMK64fxsRvs3lKc34QCXSbpk,3319
 alita_sdk/runtime/langchain/indexer.py,sha256=0ENHy5EOhThnAiYFc7QAsaTNp9rr8hDV_hTK8ahbatk,37592
 alita_sdk/runtime/langchain/langraph_agent.py,sha256=BkrbYMy4BPAvy9uANH3s9ffBzaGewKFK97evN90L5kY,39917
 alita_sdk/runtime/langchain/mixedAgentParser.py,sha256=M256lvtsL3YtYflBCEp-rWKrKtcY1dJIyRGVv7KW9ME,2611
 alita_sdk/runtime/langchain/mixedAgentRenderes.py,sha256=asBtKqm88QhZRILditjYICwFVKF5KfO38hu2O-WrSWE,5964
+alita_sdk/runtime/langchain/store_manager.py,sha256=2MSbk8Zy969R2YjA5lV6qIPoFRPZYAIUHyoK_2apoGw,1420
 alita_sdk/runtime/langchain/utils.py,sha256=Npferkn10dvdksnKzLJLBI5bNGQyVWTBwqp3vQtUqmY,6631
 alita_sdk/runtime/langchain/agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 alita_sdk/runtime/langchain/agents/xml_chat.py,sha256=Mx7PK5T97_GrFCwHHZ3JZP42S7MwtUzV0W-_8j6Amt8,6212
@@ -96,12 +97,12 @@ alita_sdk/runtime/llms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3
 alita_sdk/runtime/llms/alita.py,sha256=oAALCyrTbQ8gygJO7xzZLjOowpUnTbh8R363argRUVs,10119
 alita_sdk/runtime/llms/preloaded.py,sha256=TFdoScswWI5vMkqFwEyQn_yW3mIoebWLxjciz_Ars3w,11222
 alita_sdk/runtime/toolkits/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-alita_sdk/runtime/toolkits/application.py,sha256=LrxbBV05lkRP3_WtKGBKtMdoQHXVY-_AtFr1cUuHz40,2341
+alita_sdk/runtime/toolkits/application.py,sha256=akqUuaIL9u7-SsUmS-XgN4qxDEnXFhsK9do4n8inpSo,2432
 alita_sdk/runtime/toolkits/artifact.py,sha256=7fTr9VpGd2zwCB3EwW4aqWa5jVKRTunqV3ETMJWGNmU,2356
 alita_sdk/runtime/toolkits/datasource.py,sha256=qk78OdPoReYPCWwahfkKLbKc4pfsu-061oXRryFLP6I,2498
 alita_sdk/runtime/toolkits/prompt.py,sha256=WIpTkkVYWqIqOWR_LlSWz3ug8uO9tm5jJ7aZYdiGRn0,1192
 alita_sdk/runtime/toolkits/subgraph.py,sha256=ZYqI4yVLbEPAjCR8dpXbjbL2ipX598Hk3fL6AgaqFD4,1758
-alita_sdk/runtime/toolkits/tools.py,sha256=jUjOtv1nVkZIJpDnuOwlWMZamnpFsDH7Yw1mYPzzc4U,5462
+alita_sdk/runtime/toolkits/tools.py,sha256=7a1miPVtSQHgEfCnWuX83O-NKP71EYKYi65TupLHe20,6094
 alita_sdk/runtime/toolkits/vectorstore.py,sha256=BGppQADa1ZiLO17fC0uCACTTEvPHlodEDYEzUcBRbAA,2901
 alita_sdk/runtime/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 alita_sdk/runtime/tools/agent.py,sha256=m98QxOHwnCRTT9j18Olbb5UPS8-ZGeQaGiUyZJSyFck,3162
@@ -247,6 +248,7 @@ alita_sdk/tools/llm/llm_utils.py,sha256=v3_lWP_Nk6tJLkj0BYohOun0OWNfvzqLjPdPAMl-
 alita_sdk/tools/localgit/__init__.py,sha256=NScO0Eu-wl-rc63jjD5Qv1RXXB1qukSIJXx-yS_JQLI,2529
 alita_sdk/tools/localgit/local_git.py,sha256=gsAftNcK7nMCd8VsIkwDLs2SoG0MgpYdkQG5tmoynkA,18074
 alita_sdk/tools/localgit/tool.py,sha256=It_B24rMvFPurB355Oy5IShg2BsZTASsEoSS8hu2SXw,998
+alita_sdk/tools/memory/__init__.py,sha256=ZpCQQuyGjISCt0zZxMSkDLiGOYPlCPGRzttQGYeY1ts,1530
 alita_sdk/tools/ocr/__init__.py,sha256=pvslKVXyJmK0q23FFDNieuc7RBIuzNXTjTNj-GqhGb0,3335
 alita_sdk/tools/ocr/api_wrapper.py,sha256=08UF8wj1sR8DcW0z16pw19bgLatLkBF8dySW-Ds8iRk,29649
 alita_sdk/tools/ocr/text_detection.py,sha256=1DBxt54r3_HdEi93QynSIVta3rH3UpIvy799TPtDTtk,23825
@@ -312,8 +314,8 @@ alita_sdk/tools/zephyr_enterprise/api_wrapper.py,sha256=Ir3zHljhbZQJRJJQOBzS_GL5
 alita_sdk/tools/zephyr_enterprise/zephyr_enterprise.py,sha256=hV9LIrYfJT6oYp-ZfQR0YHflqBFPsUw2Oc55HwK0H48,6809
 alita_sdk/tools/zephyr_scale/__init__.py,sha256=2NTcdrfkx4GSegqyXhsPLsEpc4FlACuDy85b0fk6cAo,4572
 alita_sdk/tools/zephyr_scale/api_wrapper.py,sha256=UHVQUVqcBc3SZvDfO78HSuBzwAsRw2cCDQa-xMOzndE,68663
-alita_sdk-0.3.148.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-alita_sdk-0.3.148.dist-info/METADATA,sha256=4YPh627MjXCdFWepoidv1eLeFpD3iDrMBldXGeUJ-bE,18418
-alita_sdk-0.3.148.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-alita_sdk-0.3.148.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
-alita_sdk-0.3.148.dist-info/RECORD,,
+alita_sdk-0.3.150.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+alita_sdk-0.3.150.dist-info/METADATA,sha256=N2t0d8uU3QTSNjZU84fyCNzntCQlzlQoow-vd8lROyY,18667
+alita_sdk-0.3.150.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+alita_sdk-0.3.150.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
+alita_sdk-0.3.150.dist-info/RECORD,,