pedrobento988 commited on
Commit
2710472
Β·
verified Β·
1 Parent(s): aa722cf

Add project structure files (#3)

Browse files

- chore: Add repository structure files (3ea4ef03ff1d1cba8bfa04637618f0a28a659371)
- feat: Add langgraph chat with bedrock (4a76b1f710e84b6b7cfc6a05e946df474b4d1b91)

Files changed (10) hide show
  1. .gitignore +130 -0
  2. .pre-commit-config.yaml +118 -0
  3. README.md +17 -1
  4. app.py +2 -62
  5. pyproject.toml +130 -0
  6. requirements-dev.txt +55 -0
  7. requirements.txt +25 -1
  8. tdagent/__init__.py +0 -0
  9. tdagent/grchat.py +202 -0
  10. uv.lock +0 -0
.gitignore ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ .Python
11
+ env/
12
+ build/
13
+ develop-eggs/
14
+ dist/
15
+ downloads/
16
+ eggs/
17
+ .eggs/
18
+ lib/
19
+ lib64/
20
+ parts/
21
+ sdist/
22
+ var/
23
+ wheels/
24
+ *.egg-info/
25
+ .installed.cfg
26
+ *.egg
27
+
28
+ # PyInstaller
29
+ # Usually these files are written by a python script from a template
30
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
31
+ *.manifest
32
+ *.spec
33
+
34
+ # Installer logs
35
+ pip-log.txt
36
+ pip-delete-this-directory.txt
37
+
38
+ # Unit test / coverage reports
39
+ htmlcov/
40
+ .tox/
41
+ .coverage
42
+ .coverage.*
43
+ .cache
44
+ nosetests.xml
45
+ coverage.xml
46
+ *.cover
47
+ .hypothesis/
48
+ .pytest_cache/
49
+
50
+ # Translations
51
+ *.mo
52
+ *.pot
53
+
54
+ # PyBuilder
55
+ target/
56
+
57
+ # Jupyter Notebook
58
+ .ipynb_checkpoints
59
+
60
+ # pyenv
61
+ .python-version
62
+
63
+ # dotenv
64
+ .env
65
+
66
+ # virtualenv
67
+ .venv
68
+ venv/
69
+ ENV/
70
+
71
+ # Sphinx documentation
72
+ docs/_build/
73
+
74
+ # mkdocs documentation
75
+ /site
76
+ docs/*.png
77
+
78
+ # mypy
79
+ .mypy_cache/
80
+
81
+ # vscode cache
82
+ .vscode
83
+
84
+ # Pycharm files
85
+ .idea
86
+
87
+ ### macOS template
88
+ *.DS_Store
89
+ .AppleDouble
90
+ .LSOverride
91
+
92
+ # serverless
93
+ **/.serverless
94
+ **/node_modules
95
+
96
+ # terraform
97
+ **/.terraform
98
+ *.tfstate*
99
+
100
+ # exclude temp files from source control
101
+ temp/
102
+ tmp/
103
+
104
+ # exclude data from source control by default
105
+ /data/*
106
+
107
+ !/data/sat2023-sample/
108
+ !/data/cnf-sample/
109
+
110
+ !/data/sat2023-cos-similarity-dataset/
111
+ /data/sat2023-cos-similarity-dataset/*
112
+ !/data/sat2023-cos-similarity-dataset/results_main_detailed.csv
113
+
114
+ # Configurations
115
+ /config/
116
+
117
+ # Logs
118
+ lightning_logs/
119
+
120
+ # Models
121
+ /models/
122
+
123
+
124
+ # Poetry code artifact settings
125
+ .poetry/
126
+ poetry.toml
127
+
128
+ # Mise
129
+ .mise.toml
130
+ mise.toml
.pre-commit-config.yaml ADDED
@@ -0,0 +1,118 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ repos:
2
+ - repo: https://github.com/pre-commit/pre-commit-hooks
3
+ rev: v2.4.0
4
+ hooks:
5
+ - id: requirements-txt-fixer
6
+ files: requirements.txt|requirements-dev.txt|requirements-test.txt
7
+ - id: trailing-whitespace
8
+ exclude: |
9
+ (?x)^(
10
+ notebooks/
11
+ )
12
+ args: [--markdown-linebreak-ext=md]
13
+ - id: end-of-file-fixer
14
+ exclude: |
15
+ (?x)^(
16
+ notebooks/
17
+ )
18
+ - id: check-yaml
19
+ - id: check-symlinks
20
+ - id: check-toml
21
+ - id: check-added-large-files
22
+ args: ["--maxkb=1000"]
23
+ - repo: https://github.com/asottile/add-trailing-comma
24
+ rev: v3.1.0
25
+ hooks:
26
+ - id: add-trailing-comma
27
+ - repo: https://github.com/psf/black
28
+ rev: 23.1.0
29
+ hooks:
30
+ - id: black
31
+ exclude: |
32
+ (?x)^(
33
+ notebooks/
34
+ )
35
+ - repo: https://github.com/pycqa/isort
36
+ rev: "5.12.0"
37
+ hooks:
38
+ - id: isort
39
+ exclude: |
40
+ (?x)^(
41
+ notebooks/
42
+ )
43
+ - repo: https://github.com/astral-sh/ruff-pre-commit
44
+ rev: v0.9.7
45
+ hooks:
46
+ - id: ruff # linter
47
+ exclude: |
48
+ (?x)^(
49
+ scripts/|
50
+ notebooks/
51
+ )
52
+ # - id: ruff-format
53
+ - repo: local
54
+ hooks:
55
+ - id: update-req
56
+ name: Update requirements.txt
57
+ stages: [pre-commit]
58
+ language: system
59
+ entry: uv
60
+ files: uv.lock|requirements.txt
61
+ pass_filenames: false
62
+ args:
63
+ [
64
+ "export",
65
+ "--format",
66
+ "requirements-txt",
67
+ "--no-hashes",
68
+ "--no-dev",
69
+ "-o",
70
+ "requirements.txt",
71
+ ]
72
+ - id: update-dev-req
73
+ name: Update requirements-dev.txt
74
+ stages: [pre-commit]
75
+ language: system
76
+ entry: uv
77
+ files: uv.lock|requirements-dev.txt
78
+ pass_filenames: false
79
+ args:
80
+ [
81
+ "export",
82
+ "--format",
83
+ "requirements-txt",
84
+ "--no-hashes",
85
+ "--group",
86
+ "dev",
87
+ "--group",
88
+ "test",
89
+ "-o",
90
+ "requirements-dev.txt",
91
+ ]
92
+ - id: mypy
93
+ name: Running mypy
94
+ stages: [commit]
95
+ language: system
96
+ entry: uv run mypy
97
+ args: [--install-types, --non-interactive]
98
+ types: [python]
99
+ exclude: |
100
+ (?x)^(
101
+ scripts/|
102
+ notebooks/
103
+ )
104
+
105
+ # - id: pytest
106
+ # name: pytest
107
+ # stages: [commit]
108
+ # language: system
109
+ # entry: poetry run pytest
110
+ # types: [python]
111
+
112
+ # - id: pytest-cov
113
+ # name: pytest
114
+ # stages: [push]
115
+ # language: system
116
+ # entry: poetry run pytest --cov --cov-fail-under=100
117
+ # types: [python]
118
+ # pass_filenames: false
README.md CHANGED
@@ -11,4 +11,20 @@ license: apache-2.0
11
  short_description: tdb
12
  ---
13
 
14
- An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11
  short_description: tdb
12
  ---
13
 
14
+ # TDA Agent
15
+
16
+ # Development setup
17
+
18
+ To start developing you need the following tools:
19
+
20
+ * [uv](https://docs.astral.sh/uv/)
21
+
22
+ To start, sync all the dependencies with `uv sync --all-groups`.
23
+ Then, install the pre-commit hooks (`uv run pre-commit install`) to
24
+ ensure that future commits comply with the bare minimum to keep
25
+ code _readable_.
26
+
27
+
28
+ ## Old content
29
+
30
+ An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).
app.py CHANGED
@@ -1,64 +1,4 @@
1
- import gradio as gr
2
- from huggingface_hub import InferenceClient
3
-
4
- """
5
- For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
6
- """
7
- client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
8
-
9
-
10
- def respond(
11
- message,
12
- history: list[tuple[str, str]],
13
- system_message,
14
- max_tokens,
15
- temperature,
16
- top_p,
17
- ):
18
- messages = [{"role": "system", "content": system_message}]
19
-
20
- for val in history:
21
- if val[0]:
22
- messages.append({"role": "user", "content": val[0]})
23
- if val[1]:
24
- messages.append({"role": "assistant", "content": val[1]})
25
-
26
- messages.append({"role": "user", "content": message})
27
-
28
- response = ""
29
-
30
- for message in client.chat_completion(
31
- messages,
32
- max_tokens=max_tokens,
33
- stream=True,
34
- temperature=temperature,
35
- top_p=top_p,
36
- ):
37
- token = message.choices[0].delta.content
38
-
39
- response += token
40
- yield response
41
-
42
-
43
- """
44
- For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
45
- """
46
- demo = gr.ChatInterface(
47
- respond,
48
- additional_inputs=[
49
- gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
50
- gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
51
- gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
52
- gr.Slider(
53
- minimum=0.1,
54
- maximum=1.0,
55
- value=0.95,
56
- step=0.05,
57
- label="Top-p (nucleus sampling)",
58
- ),
59
- ],
60
- )
61
-
62
 
63
  if __name__ == "__main__":
64
- demo.launch()
 
1
+ from tdagent.grchat import gr_app
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
  if __name__ == "__main__":
4
+ gr_app.launch()
pyproject.toml ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [project]
2
+ name = "tdagent"
3
+ version = "0.1.0"
4
+ description = "TDA Agent implemented for huggingface hackathon."
5
+ authors = [
6
+ { name = "Pedro Completo Bento", email = "[email protected]" },
7
+ { name = "Josep Pon Farreny", email = "[email protected]" },
8
+ { name = "Miguel Rodin Rodriguez", email = "[email protected]" },
9
+ { name = "Sofia Jeronimo dos Santos", email = "[email protected]" },
10
+ { name = "Rodrigo Dominguez Sanz", email = "[email protected]" },
11
+ ]
12
+ requires-python = ">=3.10,<4"
13
+ readme = "README.md"
14
+ license = ""
15
+ dependencies = [
16
+ "gradio[mcp]>=5.32.1",
17
+ "huggingface-hub>=0.32.3",
18
+ "langchain-aws>=0.2.24",
19
+ "langgraph>=0.4.7",
20
+ ]
21
+
22
+ [project.scripts]
23
+
24
+
25
+ [dependency-groups]
26
+ dev = ["mypy~=1.14", "ruff>=0.9,<1", "pre-commit~=3.4", "pip-audit>=2.9.0"]
27
+ test = [
28
+ "pytest>=7.4.4,<8",
29
+ "pytest-cov>=4.1.0,<5",
30
+ "pytest-randomly>=3.15.0,<4",
31
+ "xdoctest>=1.1.2,<2",
32
+ ]
33
+
34
+ [build-system]
35
+ requires = ["hatchling"]
36
+ build-backend = "hatchling.build"
37
+
38
+ [tool.uv]
39
+ package = false
40
+ default-groups = ["dev", "test"]
41
+
42
+ [tool.uv.workspace]
43
+ members = ["test"]
44
+
45
+ [tool.black]
46
+ target-version = ["py39", "py310", "py311"]
47
+ line-length = 88
48
+
49
+ [tool.isort]
50
+ profile = "black"
51
+ lines_after_imports = 2
52
+
53
+ [tool.mypy]
54
+ cache_dir = ".cache/mypy/"
55
+ ignore_missing_imports = true
56
+ no_implicit_optional = true
57
+ check_untyped_defs = true
58
+ strict_equality = true
59
+ disallow_any_generics = true
60
+ disallow_subclassing_any = true
61
+ disallow_untyped_calls = true
62
+ disallow_untyped_defs = true
63
+ disallow_incomplete_defs = true
64
+ disallow_untyped_decorators = true
65
+ warn_redundant_casts = true
66
+ warn_unused_ignores = true
67
+ exclude = "docs/"
68
+ plugins = ["pydantic.mypy"] # ["numpy.typing.mypy_plugin"]
69
+
70
+ [[tool.mypy.overrides]]
71
+ module = "tests.*"
72
+ disallow_untyped_defs = false
73
+ disallow_incomplete_defs = false
74
+
75
+ [tool.pytest.ini_options]
76
+ cache_dir = ".cache"
77
+ testpaths = ["tests", "tda_agent"]
78
+ addopts = [
79
+ "--strict",
80
+ "-r sxX",
81
+ "--cov-report=html",
82
+ "--cov-report=term-missing:skip-covered",
83
+ "--no-cov-on-fail",
84
+ "--xdoc",
85
+ ]
86
+ console_output_style = "count"
87
+ markers = ""
88
+ filterwarnings = ["ignore::DeprecationWarning"]
89
+
90
+ [tool.ruff]
91
+ cache-dir = ".cache/ruff"
92
+ exclude = [
93
+ ".git",
94
+ "__pycache__",
95
+ "docs/source/conf.py",
96
+ "old",
97
+ "build",
98
+ "dist",
99
+ ".venv",
100
+ "scripts",
101
+ ]
102
+ line-length = 88
103
+
104
+ [tool.ruff.lint]
105
+ select = ["ALL"]
106
+ ignore = ["D100", "D104", "D107", "D401", "EM102", "ERA001", "TRY003"]
107
+
108
+ [tool.ruff.lint.flake8-quotes]
109
+ inline-quotes = "double"
110
+
111
+ [tool.ruff.lint.flake8-bugbear]
112
+ # Allow default arguments like, e.g., `data: List[str] = fastapi.Query(None)`.
113
+ extend-immutable-calls = ["typer.Argument", "typer.Option"]
114
+
115
+ [tool.ruff.lint.pep8-naming]
116
+ ignore-names = ["F", "L"]
117
+
118
+ [tool.ruff.lint.isort]
119
+ lines-after-imports = 2
120
+
121
+ [tool.ruff.lint.mccabe]
122
+ max-complexity = 18
123
+
124
+ [tool.ruff.lint.pydocstyle]
125
+ convention = "google"
126
+
127
+ [tool.ruff.lint.per-file-ignores]
128
+ "*/__init__.py" = ["F401"]
129
+ "tdagent/cli/**/*.py" = ["D103", "T201"]
130
+ "tests/*.py" = ["D103", "PLR2004", "S101"]
requirements-dev.txt ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # This file was autogenerated by uv via the following command:
2
+ # uv export --format requirements-txt --no-hashes --group dev --group test -o requirements-dev.txt
3
+ boolean-py==5.0
4
+ cachecontrol==0.14.3
5
+ certifi==2025.4.26
6
+ cfgv==3.4.0
7
+ charset-normalizer==3.4.2
8
+ colorama==0.4.6 ; sys_platform == 'win32'
9
+ coverage==7.8.2
10
+ cyclonedx-python-lib==9.1.0
11
+ defusedxml==0.7.1
12
+ distlib==0.3.9
13
+ exceptiongroup==1.3.0 ; python_full_version < '3.11'
14
+ filelock==3.18.0
15
+ fsspec==2025.5.1
16
+ hf-xet==1.1.2 ; platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
17
+ huggingface-hub==0.32.3
18
+ identify==2.6.12
19
+ idna==3.10
20
+ iniconfig==2.1.0
21
+ license-expression==30.4.1
22
+ markdown-it-py==3.0.0
23
+ mdurl==0.1.2
24
+ msgpack==1.1.0
25
+ mypy==1.16.0
26
+ mypy-extensions==1.1.0
27
+ nodeenv==1.9.1
28
+ packageurl-python==0.16.0
29
+ packaging==25.0
30
+ pathspec==0.12.1
31
+ pip==25.1.1
32
+ pip-api==0.0.34
33
+ pip-audit==2.9.0
34
+ pip-requirements-parser==32.0.1
35
+ platformdirs==4.3.8
36
+ pluggy==1.6.0
37
+ pre-commit==3.8.0
38
+ py-serializable==2.0.0
39
+ pygments==2.19.1
40
+ pyparsing==3.2.3
41
+ pytest==7.4.4
42
+ pytest-cov==4.1.0
43
+ pytest-randomly==3.16.0
44
+ pyyaml==6.0.2
45
+ requests==2.32.3
46
+ rich==14.0.0
47
+ ruff==0.11.12
48
+ sortedcontainers==2.4.0
49
+ toml==0.10.2
50
+ tomli==2.2.1 ; python_full_version <= '3.11'
51
+ tqdm==4.67.1
52
+ typing-extensions==4.13.2
53
+ urllib3==2.4.0
54
+ virtualenv==20.31.2
55
+ xdoctest==1.2.0
requirements.txt CHANGED
@@ -1 +1,25 @@
1
- huggingface_hub==0.25.2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # This file was autogenerated by uv via the following command:
2
+ # uv export --format requirements-txt --no-hashes --no-dev -o requirements.txt
3
+ certifi==2025.4.26
4
+ charset-normalizer==3.4.2
5
+ colorama==0.4.6 ; sys_platform == 'win32'
6
+ coverage==7.8.2
7
+ exceptiongroup==1.3.0 ; python_full_version < '3.11'
8
+ filelock==3.18.0
9
+ fsspec==2025.5.1
10
+ hf-xet==1.1.2 ; platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
11
+ huggingface-hub==0.32.3
12
+ idna==3.10
13
+ iniconfig==2.1.0
14
+ packaging==25.0
15
+ pluggy==1.6.0
16
+ pytest==7.4.4
17
+ pytest-cov==4.1.0
18
+ pytest-randomly==3.16.0
19
+ pyyaml==6.0.2
20
+ requests==2.32.3
21
+ tomli==2.2.1 ; python_full_version <= '3.11'
22
+ tqdm==4.67.1
23
+ typing-extensions==4.13.2
24
+ urllib3==2.4.0
25
+ xdoctest==1.2.0
tdagent/__init__.py ADDED
File without changes
tdagent/grchat.py ADDED
@@ -0,0 +1,202 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ from collections.abc import Mapping
4
+ from types import MappingProxyType
5
+
6
+ import boto3
7
+ import botocore
8
+ import botocore.exceptions
9
+ import gradio as gr
10
+ from langchain_aws import ChatBedrock
11
+ from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
12
+ from langgraph.graph.graph import CompiledGraph
13
+ from langgraph.prebuilt import create_react_agent
14
+
15
+
16
+ #### Constants ####
17
+
18
+ SYSTEM_MESSAGE = SystemMessage(
19
+ "You are a helpful assistant.",
20
+ )
21
+
22
+ GRADIO_ROLE_TO_LG_MESSAGE_TYPE = MappingProxyType(
23
+ {
24
+ "user": HumanMessage,
25
+ "assistant": AIMessage,
26
+ },
27
+ )
28
+
29
+
30
+ #### Shared variables ####
31
+
32
+ llm_agent: CompiledGraph | None = None
33
+
34
+ #### Utility functions ####
35
+
36
+
37
+ def create_bedrock_llm(
38
+ bedrock_model_id: str,
39
+ aws_access_key: str,
40
+ aws_secret_key: str,
41
+ aws_session_token: str,
42
+ aws_region: str,
43
+ ) -> tuple[ChatBedrock | None, str]:
44
+ """Create a LangGraph Bedrock agent."""
45
+ boto3_config = {
46
+ "aws_access_key_id": aws_access_key,
47
+ "aws_secret_access_key": aws_secret_key,
48
+ "aws_session_token": aws_session_token if aws_session_token else None,
49
+ "region_name": aws_region,
50
+ }
51
+
52
+ # Verify credentials
53
+ try:
54
+ sts = boto3.client("sts", **boto3_config)
55
+ sts.get_caller_identity()
56
+ except botocore.exceptions.ClientError as err:
57
+ return None, str(err)
58
+
59
+ try:
60
+ bedrock_client = boto3.client("bedrock-runtime", **boto3_config)
61
+ llm = ChatBedrock(
62
+ model=bedrock_model_id,
63
+ client=bedrock_client,
64
+ model_kwargs={"temperature": 0.7},
65
+ )
66
+ except Exception as e: # noqa: BLE001
67
+ return None, str(e)
68
+
69
+ return llm, ""
70
+
71
+
72
+ #### UI functionality ####
73
+
74
+
75
+ async def gr_connect_to_bedrock(
76
+ model_id: str,
77
+ access_key: str,
78
+ secret_key: str,
79
+ session_token: str,
80
+ region: str,
81
+ ) -> str:
82
+ """Initialize Bedrock agent."""
83
+ global llm_agent # noqa: PLW0603
84
+
85
+ if not access_key or not secret_key:
86
+ return "❌ Please provide both Access Key ID and Secret Access Key"
87
+
88
+ llm, error = create_bedrock_llm(
89
+ model_id,
90
+ access_key.strip(),
91
+ secret_key.strip(),
92
+ session_token.strip(),
93
+ region,
94
+ )
95
+
96
+ if llm is None:
97
+ return f"❌ Connection failed: {error}"
98
+
99
+ llm_agent = create_react_agent(
100
+ model=llm,
101
+ tools=[],
102
+ prompt=SYSTEM_MESSAGE,
103
+ )
104
+
105
+ return "βœ… Successfully connected to AWS Bedrock!"
106
+
107
+
108
+ async def gr_chat_function( # noqa: D103
109
+ message: str,
110
+ history: list[Mapping[str, str]],
111
+ ) -> str:
112
+ if llm_agent is None:
113
+ return "Please configure your credentials first."
114
+
115
+ messages = []
116
+ for hist_msg in history:
117
+ role = hist_msg["role"]
118
+ message_type = GRADIO_ROLE_TO_LG_MESSAGE_TYPE[role]
119
+ messages.append(message_type(content=hist_msg["content"]))
120
+
121
+ messages.append(HumanMessage(content=message))
122
+
123
+ llm_response = await llm_agent.ainvoke(
124
+ {
125
+ "messages": messages,
126
+ },
127
+ )
128
+
129
+ return llm_response["messages"][-1].content
130
+
131
+
132
+ ## UI components ##
133
+
134
+ with gr.Blocks() as gr_app:
135
+ gr.Markdown("# πŸ” Secure Bedrock Chatbot")
136
+
137
+ # Credentials section (collapsible)
138
+ with gr.Accordion("πŸ”‘ Bedrock Configuration", open=True):
139
+ gr.Markdown(
140
+ "**Note**: Credentials are only stored in memory during your session.",
141
+ )
142
+ with gr.Row():
143
+ bedrock_model_id_textbox = gr.Textbox(
144
+ label="Bedrock Model Id",
145
+ value="eu.anthropic.claude-3-5-sonnet-20240620-v1:0",
146
+ )
147
+ with gr.Row():
148
+ aws_access_key_textbox = gr.Textbox(
149
+ label="AWS Access Key ID",
150
+ type="password",
151
+ placeholder="Enter your AWS Access Key ID",
152
+ )
153
+ aws_secret_key_textbox = gr.Textbox(
154
+ label="AWS Secret Access Key",
155
+ type="password",
156
+ placeholder="Enter your AWS Secret Access Key",
157
+ )
158
+ with gr.Row():
159
+ aws_session_token_textbox = gr.Textbox(
160
+ label="AWS Session Token",
161
+ type="password",
162
+ placeholder="Enter your AWS session token",
163
+ )
164
+ with gr.Row():
165
+ aws_region_dropdown = gr.Dropdown(
166
+ label="AWS Region",
167
+ choices=[
168
+ "us-east-1",
169
+ "us-west-2",
170
+ "eu-west-1",
171
+ "eu-central-1",
172
+ "ap-southeast-1",
173
+ ],
174
+ value="eu-west-1",
175
+ )
176
+ connect_btn = gr.Button("πŸ”Œ Connect to Bedrock", variant="primary")
177
+
178
+ status_textbox = gr.Textbox(label="Connection Status", interactive=False)
179
+
180
+ connect_btn.click(
181
+ gr_connect_to_bedrock,
182
+ inputs=[
183
+ bedrock_model_id_textbox,
184
+ aws_access_key_textbox,
185
+ aws_secret_key_textbox,
186
+ aws_session_token_textbox,
187
+ aws_region_dropdown,
188
+ ],
189
+ outputs=[status_textbox],
190
+ )
191
+
192
+ chat_interface = gr.ChatInterface(
193
+ fn=gr_chat_function,
194
+ type="messages",
195
+ examples=[],
196
+ title="Agent with MCP Tools",
197
+ description="This is a simple agent that uses MCP tools.",
198
+ )
199
+
200
+
201
+ if __name__ == "__main__":
202
+ gr_app.launch()
uv.lock ADDED
The diff for this file is too large to render. See raw diff