llama-index-tools-azure-code-interpreter 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,105 @@
1
+ Metadata-Version: 2.1
2
+ Name: llama-index-tools-azure-code-interpreter
3
+ Version: 0.1.0
4
+ Summary: llama-index tools azure dynamic sessions integration for code interpreter
5
+ License: MIT
6
+ Author: Your Name
7
+ Author-email: you@example.com
8
+ Maintainer: ajhofmann
9
+ Requires-Python: >=3.8.1,<4.0
10
+ Classifier: License :: OSI Approved :: MIT License
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Programming Language :: Python :: 3.9
13
+ Classifier: Programming Language :: Python :: 3.10
14
+ Classifier: Programming Language :: Python :: 3.11
15
+ Requires-Dist: azure-identity (>=1.16.0,<2.0.0)
16
+ Requires-Dist: llama-index-core (>=0.10.1,<0.11.0)
17
+ Requires-Dist: requests (>=2.31.0,<3.0.0)
18
+ Description-Content-Type: text/markdown
19
+
20
+ # Azure Code Interpreter Tool
21
+
22
+ This tool leverages Azure Dynamic Sessions Pool to enable an Agent to run generated Python code in a secure environment with very low latency.
23
+
24
+ In order to utilize the tool, you will need to have the Session Pool management endpoint first. [Learn more](https://aka.ms/aca/sessions)
25
+
26
+ ## Prerequisites
27
+
28
+ - Make sure to create a Session Pool and note down the `poolManagementEndpoint`.
29
+
30
+ - In order to have the code execution right, the correct role needs to be assigned to the current user agent. Be sure to assign `Session Pool Executor` role to the correct user agent's identity (e.g. User Email, Service Principal, Managed Identity, etc.) in Session Pool's access control panel through the Portal or CLI. [Learn more](https://aka.ms/aca/sessions)
31
+
32
+ ## Usage
33
+
34
+ A more detailed sample is located in a Jupyter notebook [here](https://github.com/run-llama/llama_index/tree/main/docs/docs/examples/tools/azure_code_interpreter.ipynb)
35
+
36
+ Here's an example usage of the `AzureCodeInterpreterToolSpec`.
37
+
38
+ 1. First, install the Azure Dynamic Sessions package using `pip`:
39
+
40
+ ```
41
+ pip install llama-index-tools-azure-code-interpreter
42
+ ```
43
+
44
+ 2. Next, set up the Dynamic Sessions tool and a LLM agent:
45
+
46
+ ```python
47
+ from llama_index.tools.azure_code_interpreter import (
48
+ AzureCodeInterpreterToolSpec,
49
+ )
50
+ from llama_index.core.agent import ReActAgent
51
+ from llama_index.llms.azure_openai import AzureOpenAI
52
+
53
+ llm = AzureOpenAI(
54
+ model="gpt-35-turbo",
55
+ deployment_name="gpt-35-deploy",
56
+ api_key=api_key,
57
+ azure_endpoint=azure_endpoint,
58
+ api_version=api_version,
59
+ )
60
+
61
+ code_interpreter_spec = AzureCodeInterpreterToolSpec(
62
+ pool_managment_endpoint="your-pool-management-endpoint"
63
+ )
64
+
65
+ agent = ReActAgent.from_tools(
66
+ code_interpreter_spec.to_tool_list(), llm=llm, verbose=True
67
+ )
68
+ ```
69
+
70
+ 3. Use the tool as you need:
71
+
72
+ ```python
73
+ print(agent.chat("Tell me the current time in Seattle."))
74
+
75
+ """
76
+ Sample Return:
77
+ Thought: To provide the current time in Seattle, I need to calculate it based on the current UTC time and adjust for Seattle's time zone, which is Pacific Daylight Time (PDT) during daylight saving time and Pacific Standard Time (PST) outside of daylight saving time. PDT is UTC-7, and PST is UTC-8. I can use the code interpreter tool to get the current UTC time and adjust it accordingly.
78
+ Action: code_interpreter
79
+ Action Input: {'python_code': "from datetime import datetime, timedelta; import pytz; utc_now = datetime.now(pytz.utc); seattle_time = utc_now.astimezone(pytz.timezone('America/Los_Angeles')); seattle_time.strftime('%Y-%m-%d %H:%M:%S %Z%z')"}
80
+ Observation: {'$id': '1', 'status': 'Success', 'stdout': '', 'stderr': '', 'result': '2024-05-04 13:54:09 PDT-0700', 'executionTimeInMilliseconds': 120}
81
+ Thought: I can answer without using any more tools. I'll use the user's language to answer.
82
+ Answer: The current time in Seattle is 2024-05-04 13:54:09 PDT.
83
+ The current time in Seattle is 2024-05-04 13:54:09 PDT.
84
+ """
85
+
86
+ print(dynamic_session_tool.code_interpreter("1+1"))
87
+
88
+ """
89
+ Sample Return:
90
+ {'$id': '1', 'status': 'Success', 'stdout': '', 'stderr': '', 'result': 2, 'executionTimeInMilliseconds': 11}
91
+ """
92
+ ```
93
+
94
+ ## Included Tools
95
+
96
+ The `AzureCodeInterpreterToolSpec` provides the following tools to the agent:
97
+
98
+ `code_interpreter`: (Available to developer and LLM Agent in tool spec) Send a Python code to be executed in Azure Container Apps Dynamic Sessions and return the output in a JSON format.
99
+
100
+ `list_files`: (Available to developer and LLM Agent in tool spec) List the files available in a Session under the path `/mnt/data`.
101
+
102
+ `upload_file`: (Available to developer) Upload a file or a stream of data into a Session under the path `/mnt/data`.
103
+
104
+ `download_file`: (Available to developer) Download a file by its path relative to the path `/mnt/data` to the tool's hosting agent.
105
+
@@ -0,0 +1,85 @@
1
+ # Azure Code Interpreter Tool
2
+
3
+ This tool leverages Azure Dynamic Sessions Pool to enable an Agent to run generated Python code in a secure environment with very low latency.
4
+
5
+ In order to utilize the tool, you will need to have the Session Pool management endpoint first. [Learn more](https://aka.ms/aca/sessions)
6
+
7
+ ## Prerequisites
8
+
9
+ - Make sure to create a Session Pool and note down the `poolManagementEndpoint`.
10
+
11
+ - In order to have the code execution right, the correct role needs to be assigned to the current user agent. Be sure to assign `Session Pool Executor` role to the correct user agent's identity (e.g. User Email, Service Principal, Managed Identity, etc.) in Session Pool's access control panel through the Portal or CLI. [Learn more](https://aka.ms/aca/sessions)
12
+
13
+ ## Usage
14
+
15
+ A more detailed sample is located in a Jupyter notebook [here](https://github.com/run-llama/llama_index/tree/main/docs/docs/examples/tools/azure_code_interpreter.ipynb)
16
+
17
+ Here's an example usage of the `AzureCodeInterpreterToolSpec`.
18
+
19
+ 1. First, install the Azure Dynamic Sessions package using `pip`:
20
+
21
+ ```
22
+ pip install llama-index-tools-azure-code-interpreter
23
+ ```
24
+
25
+ 2. Next, set up the Dynamic Sessions tool and a LLM agent:
26
+
27
+ ```python
28
+ from llama_index.tools.azure_code_interpreter import (
29
+ AzureCodeInterpreterToolSpec,
30
+ )
31
+ from llama_index.core.agent import ReActAgent
32
+ from llama_index.llms.azure_openai import AzureOpenAI
33
+
34
+ llm = AzureOpenAI(
35
+ model="gpt-35-turbo",
36
+ deployment_name="gpt-35-deploy",
37
+ api_key=api_key,
38
+ azure_endpoint=azure_endpoint,
39
+ api_version=api_version,
40
+ )
41
+
42
+ code_interpreter_spec = AzureCodeInterpreterToolSpec(
43
+ pool_managment_endpoint="your-pool-management-endpoint"
44
+ )
45
+
46
+ agent = ReActAgent.from_tools(
47
+ code_interpreter_spec.to_tool_list(), llm=llm, verbose=True
48
+ )
49
+ ```
50
+
51
+ 3. Use the tool as you need:
52
+
53
+ ```python
54
+ print(agent.chat("Tell me the current time in Seattle."))
55
+
56
+ """
57
+ Sample Return:
58
+ Thought: To provide the current time in Seattle, I need to calculate it based on the current UTC time and adjust for Seattle's time zone, which is Pacific Daylight Time (PDT) during daylight saving time and Pacific Standard Time (PST) outside of daylight saving time. PDT is UTC-7, and PST is UTC-8. I can use the code interpreter tool to get the current UTC time and adjust it accordingly.
59
+ Action: code_interpreter
60
+ Action Input: {'python_code': "from datetime import datetime, timedelta; import pytz; utc_now = datetime.now(pytz.utc); seattle_time = utc_now.astimezone(pytz.timezone('America/Los_Angeles')); seattle_time.strftime('%Y-%m-%d %H:%M:%S %Z%z')"}
61
+ Observation: {'$id': '1', 'status': 'Success', 'stdout': '', 'stderr': '', 'result': '2024-05-04 13:54:09 PDT-0700', 'executionTimeInMilliseconds': 120}
62
+ Thought: I can answer without using any more tools. I'll use the user's language to answer.
63
+ Answer: The current time in Seattle is 2024-05-04 13:54:09 PDT.
64
+ The current time in Seattle is 2024-05-04 13:54:09 PDT.
65
+ """
66
+
67
+ print(dynamic_session_tool.code_interpreter("1+1"))
68
+
69
+ """
70
+ Sample Return:
71
+ {'$id': '1', 'status': 'Success', 'stdout': '', 'stderr': '', 'result': 2, 'executionTimeInMilliseconds': 11}
72
+ """
73
+ ```
74
+
75
+ ## Included Tools
76
+
77
+ The `AzureCodeInterpreterToolSpec` provides the following tools to the agent:
78
+
79
+ `code_interpreter`: (Available to developer and LLM Agent in tool spec) Send a Python code to be executed in Azure Container Apps Dynamic Sessions and return the output in a JSON format.
80
+
81
+ `list_files`: (Available to developer and LLM Agent in tool spec) List the files available in a Session under the path `/mnt/data`.
82
+
83
+ `upload_file`: (Available to developer) Upload a file or a stream of data into a Session under the path `/mnt/data`.
84
+
85
+ `download_file`: (Available to developer) Download a file by its path relative to the path `/mnt/data` to the tool's hosting agent.
@@ -0,0 +1,4 @@
1
+ ## init
2
+ from llama_index.tools.azure_code_interpreter.base import AzureCodeInterpreterToolSpec
3
+
4
+ __all__ = ["AzureCodeInterpreterToolSpec"]
@@ -0,0 +1,286 @@
1
+ """Azure Dynamic Sessions tool spec."""
2
+
3
+ import base64
4
+ from dataclasses import dataclass
5
+ from datetime import datetime, timedelta, timezone
6
+ import importlib.metadata
7
+ from io import BufferedReader, BytesIO
8
+ from typing import Any, Callable, List, Optional
9
+ import urllib
10
+ from uuid import uuid4
11
+ import re
12
+ import os
13
+
14
+ from azure.identity import DefaultAzureCredential
15
+ from azure.core.credentials import AccessToken
16
+ import requests
17
+
18
+ from llama_index.core.tools.tool_spec.base import BaseToolSpec
19
+
20
+
21
+ @dataclass
22
+ class RemoteFileMetadata:
23
+ """Metadata for a file in the session."""
24
+
25
+ filename: str
26
+ """The filename relative to `/mnt/data`."""
27
+
28
+ size_in_bytes: int
29
+ """The size of the file in bytes."""
30
+
31
+ file_full_path: str
32
+ """The full path of the file."""
33
+
34
+ @staticmethod
35
+ def from_dict(data: dict) -> "RemoteFileMetadata":
36
+ """Create a RemoteFileMetadata object from a dictionary."""
37
+ filename = data["filename"]
38
+ if filename.startswith("/mnt/data"):
39
+ file_full_path = filename
40
+ else:
41
+ file_full_path = f"/mnt/data/{filename}"
42
+ return RemoteFileMetadata(
43
+ filename=filename,
44
+ size_in_bytes=data["size"],
45
+ file_full_path=file_full_path,
46
+ )
47
+
48
+
49
+ def _sanitize_input(query: str) -> str:
50
+ """Sanitize input and remove whitespace, backtick, and markdown.
51
+
52
+ Args:
53
+ query: The query to sanitize
54
+
55
+ Returns:
56
+ str: The sanitized query
57
+ """
58
+ # Removes `, whitespace & python from start
59
+ query = re.sub(r"^(\s|`)*(?i:python)?\s*", "", query)
60
+ # Removes whitespace & ` from end
61
+ query = re.sub(r"(\s|`)*$", "", query)
62
+ # Add new line if no new line was appended at the end of the query
63
+ if not query.endswith("\n"):
64
+ query += "\n"
65
+ return query
66
+
67
+
68
+ class AzureCodeInterpreterToolSpec(BaseToolSpec):
69
+ """Azure Code Interpreter tool spec.
70
+
71
+ Leverages Azure Dynamic Sessions to execute Python code.
72
+ """
73
+
74
+ spec_functions = ["code_interpreter", "list_files"]
75
+
76
+ def __init__(
77
+ self,
78
+ pool_managment_endpoint: str,
79
+ session_id: Optional[str] = None,
80
+ local_save_path: Optional[str] = None,
81
+ sanitize_input: bool = True,
82
+ ) -> None:
83
+ """Initialize with parameters."""
84
+ self.pool_management_endpoint: str = pool_managment_endpoint
85
+ self.access_token: Optional[AccessToken] = None
86
+
87
+ def _access_token_provider_factory() -> Callable[[], Optional[str]]:
88
+ def access_token_provider() -> Optional[str]:
89
+ """Create a function that returns an access token."""
90
+ if self.access_token is None or datetime.fromtimestamp(
91
+ self.access_token.expires_on, timezone.utc
92
+ ) < (datetime.now(timezone.utc) + timedelta(minutes=5)):
93
+ credential = DefaultAzureCredential()
94
+ self.access_token = credential.get_token(
95
+ "https://dynamicsessions.io/.default"
96
+ )
97
+ return self.access_token.token
98
+
99
+ return access_token_provider
100
+
101
+ self.access_token_provider: Callable[
102
+ [], Optional[str]
103
+ ] = _access_token_provider_factory()
104
+ """A function that returns the access token to use for the session pool."""
105
+
106
+ self.session_id: str = session_id or str(uuid4())
107
+ """The session ID to use for the session pool. Defaults to a random UUID."""
108
+
109
+ self.sanitize_input: bool = sanitize_input
110
+ """Whether to sanitize input before executing it."""
111
+
112
+ if local_save_path:
113
+ if not os.path.exists(local_save_path):
114
+ raise Exception(f"Local save path {local_save_path} does not exist.")
115
+
116
+ self.local_save_path: Optional[str] = local_save_path
117
+ """The local path to save files generated by Python interpreter."""
118
+
119
+ try:
120
+ _package_version = importlib.metadata.version(
121
+ "llamaindex-azure-code-interpreter"
122
+ )
123
+ except importlib.metadata.PackageNotFoundError:
124
+ _package_version = "0.0.0"
125
+
126
+ self.user_agent = (
127
+ f"llamaindex-azure-code-interpreter/{_package_version} (Language=Python)"
128
+ )
129
+
130
+ def _build_url(self, path: str) -> str:
131
+ pool_management_endpoint = self.pool_management_endpoint
132
+ if not pool_management_endpoint:
133
+ raise ValueError("pool_management_endpoint is not set")
134
+
135
+ if not pool_management_endpoint.endswith("/"):
136
+ pool_management_endpoint += "/"
137
+
138
+ encoded_session_id = urllib.parse.quote(self.session_id)
139
+ query = f"identifier={encoded_session_id}&api-version=2024-02-02-preview"
140
+ query_separator = "&" if "?" in pool_management_endpoint else "?"
141
+
142
+ return pool_management_endpoint + path + query_separator + query
143
+
144
+ def code_interpreter(self, python_code: str) -> dict:
145
+ """
146
+ This tool is used to execute python commands when you need to perform calculations or computations in a Session.
147
+ Input should be a valid python command. The tool returns the result, stdout, and stderr.
148
+
149
+ Args:
150
+ python_code (str): Python code to be executed generated by llm.
151
+ """
152
+ if self.sanitize_input:
153
+ python_code = _sanitize_input(python_code)
154
+
155
+ access_token = self.access_token_provider(self)
156
+ api_url = self._build_url("code/execute")
157
+ headers = {
158
+ "Authorization": f"Bearer {access_token}",
159
+ "Content-Type": "application/json",
160
+ "User-Agent": self.user_agent,
161
+ }
162
+ body = {
163
+ "properties": {
164
+ "codeInputType": "inline",
165
+ "executionType": "synchronous",
166
+ "code": python_code,
167
+ }
168
+ }
169
+
170
+ response = requests.post(api_url, headers=headers, json=body)
171
+ response.raise_for_status()
172
+ response_json = response.json()
173
+ if "result" in response_json and response_json["result"]:
174
+ if isinstance(response_json["result"], dict):
175
+ if "base64_data" in response_json["result"]:
176
+ base64_encoded_data = response_json["result"]["base64_data"]
177
+ if self.local_save_path:
178
+ file_path = f"{self.local_save_path}/{self.session_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.{response_json['result']['format']}"
179
+ decoded_data = base64.b64decode(base64_encoded_data)
180
+ with open(file_path, "wb") as f:
181
+ f.write(decoded_data)
182
+ # Check if file is written to the file path successfully. if so, update the response_json
183
+ response_json["result"]["saved_to_local_path"] = response_json[
184
+ "result"
185
+ ].pop("base64_data")
186
+ if os.path.exists(file_path):
187
+ response_json["result"]["saved_to_local_path"] = True
188
+ else:
189
+ response_json["result"]["saved_to_local_path"] = False
190
+ else:
191
+ response_json["result"]["base64_data"] = ""
192
+ return response_json
193
+
194
+ def upload_file(
195
+ self,
196
+ data: Optional[Any] = None,
197
+ local_file_path: Optional[str] = None,
198
+ ) -> List[RemoteFileMetadata]:
199
+ """Upload a file to the session under the path /mnt/data.
200
+
201
+ Args:
202
+ data: The data to upload.
203
+ local_file_path: The path to the local file to upload.
204
+
205
+ Returns:
206
+ List[RemoteFileMetadata]: The list of metadatas for the uploaded files.
207
+ """
208
+ if data and local_file_path:
209
+ raise ValueError("data and local_file_path cannot be provided together")
210
+
211
+ if local_file_path:
212
+ remote_file_path = f"/mnt/data/{os.path.basename(local_file_path)}"
213
+ data = open(local_file_path, "rb")
214
+
215
+ access_token = self.access_token_provider(self)
216
+ if not remote_file_path.startswith("/mnt/data"):
217
+ remote_file_path = f"/mnt/data/{remote_file_path}"
218
+ api_url = self._build_url("files/upload")
219
+ headers = {
220
+ "Authorization": f"Bearer {access_token}",
221
+ }
222
+
223
+ files = [("file", (remote_file_path, data, "application/octet-stream"))]
224
+
225
+ response = requests.request("POST", api_url, headers=headers, files=files)
226
+ response.raise_for_status()
227
+
228
+ response_json = response.json()
229
+ remote_files_metadatas = []
230
+ for entry in response_json["value"]:
231
+ if "properties" in entry:
232
+ remote_files_metadatas.append(
233
+ RemoteFileMetadata.from_dict(entry["properties"])
234
+ )
235
+ return remote_files_metadatas
236
+
237
+ def download_file_to_local(
238
+ self, remote_file_path: str, local_file_path: Optional[str] = None
239
+ ) -> Optional[BufferedReader]:
240
+ """Download a file from the session back to your local environment.
241
+
242
+ Args:
243
+ remote_file_path: The path to download the file from, relative to `/mnt/data`.
244
+ local_file_path: The path to save the downloaded file to. If not provided, the file is returned as a BufferedReader.
245
+
246
+ Returns:
247
+ BufferedReader: The data of the downloaded file.
248
+ """
249
+ access_token = self.access_token_provider()
250
+ # In case if the file path LLM provides is absolute, remove the /mnt/data/ prefix
251
+ remote_file_path = remote_file_path.replace("/mnt/data/", "")
252
+ api_url = self._build_url(f"files/content/{remote_file_path}")
253
+ headers = {
254
+ "Authorization": f"Bearer {access_token}",
255
+ }
256
+
257
+ response = requests.get(api_url, headers=headers)
258
+ response.raise_for_status()
259
+
260
+ if local_file_path:
261
+ with open(local_file_path, "wb") as f:
262
+ f.write(response.content)
263
+ return None
264
+
265
+ return BytesIO(response.content)
266
+
267
+ def list_files(self) -> List[RemoteFileMetadata]:
268
+ """List the files in the session.
269
+
270
+ Returns:
271
+ List[RemoteFileMetadata]: The metadata for the files in the session
272
+ """
273
+ access_token = self.access_token_provider(self)
274
+ api_url = self._build_url("files")
275
+ headers = {
276
+ "Authorization": f"Bearer {access_token}",
277
+ }
278
+
279
+ response = requests.get(api_url, headers=headers)
280
+ response.raise_for_status()
281
+
282
+ response_json = response.json()
283
+ return [
284
+ RemoteFileMetadata.from_dict(entry["properties"])
285
+ for entry in response_json["value"]
286
+ ]
@@ -0,0 +1,65 @@
1
+ [build-system]
2
+ build-backend = "poetry.core.masonry.api"
3
+ requires = ["poetry-core"]
4
+
5
+ [tool.codespell]
6
+ check-filenames = true
7
+ check-hidden = true
8
+ skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
9
+
10
+ [tool.llamahub]
11
+ contains_example = false
12
+ import_path = "llama_index.tools.azure_code_interpreter"
13
+
14
+ [tool.llamahub.class_authors]
15
+ AzureCodeInterpreterToolSpec = "harryli0108"
16
+
17
+ [tool.mypy]
18
+ disallow_untyped_defs = true
19
+ exclude = ["_static", "build", "examples", "notebooks", "venv"]
20
+ ignore_missing_imports = true
21
+ python_version = "3.8"
22
+
23
+ [tool.poetry]
24
+ authors = ["Your Name <you@example.com>"]
25
+ description = "llama-index tools azure dynamic sessions integration for code interpreter"
26
+ exclude = ["**/BUILD"]
27
+ license = "MIT"
28
+ maintainers = ["ajhofmann"]
29
+ name = "llama-index-tools-azure-code-interpreter"
30
+ readme = "README.md"
31
+ version = "0.1.0"
32
+
33
+ [tool.poetry.dependencies]
34
+ python = ">=3.8.1,<4.0"
35
+ llama-index-core = "^0.10.1"
36
+ azure-identity = "^1.16.0"
37
+ requests = "^2.31.0"
38
+
39
+ [tool.poetry.group.dev.dependencies]
40
+ ipython = "8.10.0"
41
+ jupyter = "^1.0.0"
42
+ mypy = "0.991"
43
+ pre-commit = "3.2.0"
44
+ pylint = "2.15.10"
45
+ pytest = "7.2.1"
46
+ pytest-mock = "3.11.1"
47
+ ruff = "0.0.292"
48
+ tree-sitter-languages = "^1.8.0"
49
+ types-Deprecated = ">=0.1.0"
50
+ types-PyYAML = "^6.0.12.12"
51
+ types-protobuf = "^4.24.0.4"
52
+ types-redis = "4.5.5.0"
53
+ types-requests = "2.28.11.8"
54
+ types-setuptools = "67.1.0.0"
55
+
56
+ [tool.poetry.group.dev.dependencies.black]
57
+ extras = ["jupyter"]
58
+ version = "<=23.9.1,>=23.7.0"
59
+
60
+ [tool.poetry.group.dev.dependencies.codespell]
61
+ extras = ["toml"]
62
+ version = ">=v2.2.6"
63
+
64
+ [[tool.poetry.packages]]
65
+ include = "llama_index/"